use anyhow::{Context, Result};
use rustc_hash::{FxHashMap, FxHashSet};
use std::fs::File;
use std::io::{BufRead, BufReader, BufWriter, Read, Write};
use std::path::{Path, PathBuf};
use std::process::Command;
use std::sync::atomic::{AtomicUsize, AtomicBool, Ordering};
use std::sync::{Arc, Mutex, mpsc};
use std::time::Duration;
use std::thread;
const NCBI_FTP_BASE: &str = "https://ftp.ncbi.nlm.nih.gov/genomes";
const NCBI_TAXDUMP_URL: &str = "https://ftp.ncbi.nlm.nih.gov/pub/taxonomy/taxdump.tar.gz";
const NCBI_DATASETS_API: &str = "https://api.ncbi.nlm.nih.gov/datasets/v2";
const PLSDB_META_URL: &str = "https://ccb-microbe.cs.uni-saarland.de/plsdb2025/download_meta.tar.gz";
const PLSDB_FASTA_URL: &str = "https://ccb-microbe.cs.uni-saarland.de/plsdb2025/download_fasta";
const API_BATCH_SIZE: usize = 1000;
#[derive(Clone, Debug, Default)]
pub struct PlsdbOptions {
pub dir: Option<PathBuf>,
pub skip: bool,
}
#[derive(Clone, Debug)]
pub struct FlankBuildConfig {
pub flanking_length: usize,
pub queue_buffer_gb: u32,
pub plsdb: PlsdbOptions,
}
impl Default for FlankBuildConfig {
fn default() -> Self {
Self {
flanking_length: 1000,
queue_buffer_gb: 30,
plsdb: PlsdbOptions::default(),
}
}
}
#[derive(Clone, Debug)]
pub struct AssemblyInfo {
pub accession: String,
pub taxid: String,
pub species_taxid: String,
pub organism_name: String,
}
#[derive(Clone, Debug)]
pub struct PlasmidInfo {
pub accession: String,
pub taxonomy_uid: String,
pub genus: String,
pub species: String,
}
#[derive(Clone, Debug)]
struct PafHit {
query_name: String,
query_start: usize,
query_end: usize,
gene_name: String,
gene_length: usize,
score: i64,
mapq: u8,
divergence: f32,
gap_count: usize,
raw_line: String,
}
impl PafHit {
fn from_paf_line(line: &str) -> Option<Self> {
let fields: Vec<&str> = line.split('\t').collect();
if fields.len() < 12 {
return None;
}
let query_name = fields[0].to_string();
let query_start: usize = fields[2].parse().ok()?;
let query_end: usize = fields[3].parse().ok()?;
let gene_name = fields[5].to_string();
let gene_length: usize = fields[6].parse().ok()?;
let mapq: u8 = fields[11].parse().unwrap_or(0);
let matches: usize = fields[9].parse().unwrap_or(0);
let block_len: usize = fields[10].parse().unwrap_or(0);
let mut score: i64 = 0;
let mut divergence: f32 = 1.0;
for field in &fields[12..] {
if let Some(val) = field.strip_prefix("AS:i:") {
score = val.parse().unwrap_or(0);
} else if let Some(val) = field.strip_prefix("de:f:") {
divergence = val.parse().unwrap_or(1.0);
}
}
let gap_count = block_len.saturating_sub(matches);
Some(PafHit {
query_name,
query_start,
query_end,
gene_name,
gene_length,
score,
mapq,
divergence,
gap_count,
raw_line: line.to_string(),
})
}
fn overlaps(&self, other: &PafHit) -> bool {
if self.query_name != other.query_name {
return false;
}
let start = self.query_start.max(other.query_start);
let end = self.query_end.min(other.query_end);
if start >= end {
return false;
}
let overlap = end - start;
let self_len = self.query_end - self.query_start;
let other_len = other.query_end - other.query_start;
overlap * 2 > self_len || overlap * 2 > other_len
}
}
fn compare_paf_hits(a: &PafHit, b: &PafHit) -> std::cmp::Ordering {
use std::cmp::Ordering;
b.score.cmp(&a.score)
.then_with(|| b.gene_length.cmp(&a.gene_length))
.then_with(|| b.mapq.cmp(&a.mapq))
.then_with(|| a.divergence.partial_cmp(&b.divergence).unwrap_or(Ordering::Equal))
.then_with(|| a.gap_count.cmp(&b.gap_count))
.then_with(|| a.gene_name.cmp(&b.gene_name))
}
fn deduplicate_paf_hits(hits: Vec<PafHit>) -> Vec<PafHit> {
if hits.is_empty() {
return hits;
}
let mut groups: Vec<Vec<PafHit>> = Vec::new();
for hit in hits {
let mut found = false;
for group in &mut groups {
if group.iter().any(|h| h.overlaps(&hit)) {
group.push(hit.clone());
found = true;
break;
}
}
if !found {
groups.push(vec![hit]);
}
}
groups
.into_iter()
.map(|mut g| {
g.sort_by(compare_paf_hits);
g.into_iter().next().unwrap()
})
.collect()
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum GenomeSizeCategory {
Small,
Medium,
Large,
}
const SMALL_GENOME_THRESHOLD: u64 = 5 * 1024 * 1024; const LARGE_GENOME_THRESHOLD: u64 = 20 * 1024 * 1024;
impl GenomeSizeCategory {
pub fn from_size(size_bytes: u64) -> Self {
if size_bytes < SMALL_GENOME_THRESHOLD {
GenomeSizeCategory::Small
} else if size_bytes < LARGE_GENOME_THRESHOLD {
GenomeSizeCategory::Medium
} else {
GenomeSizeCategory::Large
}
}
pub fn thread_count(&self, max_threads: usize) -> usize {
let recommended = match self {
GenomeSizeCategory::Small => 4,
GenomeSizeCategory::Medium => 6,
GenomeSizeCategory::Large => 8,
};
recommended.min(max_threads)
}
pub fn name(&self) -> &'static str {
match self {
GenomeSizeCategory::Small => "small (<5MB)",
GenomeSizeCategory::Medium => "medium (5-20MB)",
GenomeSizeCategory::Large => "large (>20MB)",
}
}
}
fn bucket_genomes_by_size(genome_files: &[PathBuf]) -> FxHashMap<GenomeSizeCategory, Vec<PathBuf>> {
let mut buckets: FxHashMap<GenomeSizeCategory, Vec<PathBuf>> = FxHashMap::default();
for path in genome_files {
let size = std::fs::metadata(path)
.map(|m| m.len())
.unwrap_or(0);
let category = GenomeSizeCategory::from_size(size);
buckets.entry(category).or_default().push(path.clone());
}
buckets
}
#[derive(Clone, Debug)]
pub struct CatalogEntry {
pub accession: String, pub taxid: String, pub species_taxid: String, pub genus: String, pub species: String, pub organism_name: String, pub source: String, }
pub struct GenomeCatalog {
entries: FxHashMap<String, CatalogEntry>,
}
impl Default for GenomeCatalog {
fn default() -> Self {
Self::new()
}
}
impl GenomeCatalog {
pub fn new() -> Self {
Self {
entries: FxHashMap::default(),
}
}
pub fn add_assembly_with_taxonomy(
&mut self,
asm: &AssemblyInfo,
source: &str,
taxonomy: Option<&TaxonomyDB>,
) {
let (parsed_genus, species) = parse_organism_name(&asm.organism_name);
let genus = resolve_genus(&asm.organism_name, &asm.taxid, &parsed_genus, taxonomy);
let entry = CatalogEntry {
accession: asm.accession.clone(),
taxid: asm.taxid.clone(),
species_taxid: asm.species_taxid.clone(),
genus,
species,
organism_name: asm.organism_name.clone(),
source: source.to_string(),
};
self.entries.insert(asm.accession.clone(), entry.clone());
if let Some(base) = asm.accession.split('.').next() {
self.entries.insert(base.to_string(), entry);
}
}
pub fn add_plasmid_with_taxonomy(
&mut self,
plasmid: &PlasmidInfo,
taxonomy: Option<&TaxonomyDB>,
) {
let organism_name = format!("{} {}", plasmid.genus, plasmid.species);
let genus = resolve_genus(&organism_name, &plasmid.taxonomy_uid, &plasmid.genus, taxonomy);
let entry = CatalogEntry {
accession: plasmid.accession.clone(),
taxid: plasmid.taxonomy_uid.clone(),
species_taxid: String::new(),
genus,
species: plasmid.species.clone(),
organism_name,
source: "plsdb".to_string(),
};
self.entries.insert(plasmid.accession.clone(), entry.clone());
let base = plasmid.accession
.strip_prefix("NZ_")
.unwrap_or(&plasmid.accession);
let no_ver = base.split('.').next().unwrap_or(base);
self.entries.insert(no_ver.to_string(), entry.clone());
let underscore_key = plasmid.accession.replace('.', "_");
self.entries.insert(underscore_key, entry);
}
pub fn get_genus(&self, accession: &str) -> Option<&str> {
if let Some(entry) = self.entries.get(accession) {
return Some(&entry.genus);
}
if let Some(base) = accession.split('.').next() {
if let Some(entry) = self.entries.get(base) {
return Some(&entry.genus);
}
}
let stripped = accession.strip_prefix("NZ_").unwrap_or(accession);
if let Some(entry) = self.entries.get(stripped) {
return Some(&entry.genus);
}
let stripped_base = stripped.split('.').next().unwrap_or(stripped);
if let Some(entry) = self.entries.get(stripped_base) {
return Some(&entry.genus);
}
if let Some(last_underscore) = accession.rfind('_') {
let suffix = &accession[last_underscore + 1..];
if !suffix.is_empty() && suffix.chars().all(|c| c.is_ascii_digit()) {
let with_dot = format!("{}.{}", &accession[..last_underscore], suffix);
if let Some(entry) = self.entries.get(&with_dot) {
return Some(&entry.genus);
}
let stripped_with_dot = with_dot.strip_prefix("NZ_").unwrap_or(&with_dot);
if let Some(entry) = self.entries.get(stripped_with_dot) {
return Some(&entry.genus);
}
}
}
None
}
pub fn save(&self, path: &Path) -> Result<()> {
let file = File::create(path)?;
let mut writer = BufWriter::new(file);
writeln!(writer, "accession\ttaxid\tspecies_taxid\tgenus\tspecies\torganism_name\tsource")?;
let mut seen: FxHashSet<&str> = FxHashSet::default();
for entry in self.entries.values() {
if seen.contains(entry.accession.as_str()) {
continue;
}
seen.insert(&entry.accession);
writeln!(
writer,
"{}\t{}\t{}\t{}\t{}\t{}\t{}",
entry.accession,
entry.taxid,
entry.species_taxid,
entry.genus,
entry.species,
entry.organism_name.replace('\t', " "),
entry.source
)?;
}
Ok(())
}
pub fn len(&self) -> usize {
let mut seen: FxHashSet<&str> = FxHashSet::default();
for entry in self.entries.values() {
seen.insert(&entry.accession);
}
seen.len()
}
#[allow(dead_code)]
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
}
fn parse_organism_name(organism: &str) -> (String, String) {
let parts: Vec<&str> = organism.split_whitespace().collect();
let genus = parts.first()
.map(|s| s.to_string())
.unwrap_or_else(|| "Unknown".to_string());
let species = if parts.len() >= 2 {
parts[1..].join(" ")
} else {
"Unknown".to_string()
};
(clean_genus(&genus), species)
}
fn clean_genus(genus: &str) -> String {
let cleaned = if let Some(idx) = genus.find('(') {
genus[..idx].trim()
} else {
genus.trim()
};
cleaned.to_string()
}
fn resolve_genus(
organism_name: &str,
taxid_str: &str,
_parsed_genus: &str,
taxonomy: Option<&TaxonomyDB>,
) -> String {
let organism_lower = organism_name.to_lowercase();
if organism_lower.contains("uncultured") {
return "uncultured".to_string();
}
if let Some(tax_db) = taxonomy {
if let Ok(taxid) = taxid_str.parse::<u32>() {
if let Some(genus) = tax_db.get_genus(taxid) {
return clean_genus(&genus);
}
if let Some((name, _rank)) = tax_db.get_genus_or_higher(taxid) {
return clean_genus(&name);
}
}
}
let cleaned = clean_genus(organism_name);
if cleaned.is_empty() {
"unknown".to_string()
} else {
cleaned
}
}
pub struct TaxonomyDB {
names: FxHashMap<u32, String>,
nodes: FxHashMap<u32, (u32, String)>,
}
impl Default for TaxonomyDB {
fn default() -> Self {
Self::new()
}
}
impl TaxonomyDB {
pub fn new() -> Self {
Self {
names: FxHashMap::default(),
nodes: FxHashMap::default(),
}
}
pub fn load(taxdump_dir: &Path) -> Result<Self> {
let mut db = Self::new();
let names_path = taxdump_dir.join("names.dmp");
let nodes_path = taxdump_dir.join("nodes.dmp");
if !names_path.exists() || !nodes_path.exists() {
anyhow::bail!("Taxonomy files not found in {}", taxdump_dir.display());
}
eprintln!(" Loading names.dmp...");
let file = File::open(&names_path)?;
let reader = BufReader::new(file);
for line in reader.lines() {
let line = line?;
let fields: Vec<&str> = line.split("\t|\t").collect();
if fields.len() < 4 {
continue;
}
let name_class = fields[3].trim_end_matches("\t|");
if name_class != "scientific name" {
continue;
}
let taxid: u32 = match fields[0].parse() {
Ok(id) => id,
Err(_) => continue,
};
let name = fields[1].to_string();
db.names.insert(taxid, name);
}
eprintln!(" Loading nodes.dmp...");
let file = File::open(&nodes_path)?;
let reader = BufReader::new(file);
for line in reader.lines() {
let line = line?;
let fields: Vec<&str> = line.split("\t|\t").collect();
if fields.len() < 3 {
continue;
}
let taxid: u32 = match fields[0].parse() {
Ok(id) => id,
Err(_) => continue,
};
let parent_taxid: u32 = match fields[1].parse() {
Ok(id) => id,
Err(_) => continue,
};
let rank = fields[2].to_string();
db.nodes.insert(taxid, (parent_taxid, rank));
}
eprintln!(" Loaded {} taxa, {} nodes", db.names.len(), db.nodes.len());
Ok(db)
}
pub fn get_genus(&self, taxid: u32) -> Option<String> {
let mut current = taxid;
let mut visited = 0;
while visited < 50 {
if let Some((parent, rank)) = self.nodes.get(¤t) {
if rank == "genus" {
return self.names.get(¤t).cloned();
}
if *parent == current {
return None;
}
current = *parent;
visited += 1;
} else {
return None;
}
}
None
}
pub fn get_genus_or_higher(&self, taxid: u32) -> Option<(String, String)> {
let mut current = taxid;
let mut visited = 0;
let target_ranks = ["genus", "family", "order", "class", "phylum"];
let mut best_match: Option<(String, String, usize)> = None;
while visited < 50 {
if let Some((parent, rank)) = self.nodes.get(¤t) {
if let Some(priority) = target_ranks.iter().position(|r| r == rank) {
if let Some(name) = self.names.get(¤t) {
if rank == "genus" {
return Some((name.clone(), rank.clone()));
}
if best_match.is_none() || priority < best_match.as_ref().unwrap().2 {
best_match = Some((name.clone(), rank.clone(), priority));
}
}
}
if *parent == current {
break;
}
current = *parent;
visited += 1;
} else {
break;
}
}
best_match.map(|(name, rank, _)| (name, rank))
}
}
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub struct BuildState {
pub started_at: String,
pub last_updated: String,
pub current_step: String,
pub completed_steps: Vec<String>,
pub genomes_downloaded: usize,
pub plsdb_extracted: usize,
pub alignments_done: bool,
pub flanking_extracted: bool,
}
impl Default for BuildState {
fn default() -> Self {
Self::new()
}
}
impl BuildState {
pub fn new() -> Self {
let now = chrono::Local::now().format("%Y-%m-%d %H:%M:%S").to_string();
Self {
started_at: now.clone(),
last_updated: now,
current_step: String::new(),
completed_steps: Vec::new(),
genomes_downloaded: 0,
plsdb_extracted: 0,
alignments_done: false,
flanking_extracted: false,
}
}
pub fn load(path: &Path) -> Option<Self> {
std::fs::read_to_string(path)
.ok()
.and_then(|s| serde_json::from_str(&s).ok())
}
pub fn save(&self, path: &Path) -> Result<()> {
let json = serde_json::to_string_pretty(self)?;
std::fs::write(path, json)?;
Ok(())
}
pub fn update_step(&mut self, step: &str) {
self.current_step = step.to_string();
self.last_updated = chrono::Local::now().format("%Y-%m-%d %H:%M:%S").to_string();
}
pub fn complete_step(&mut self, step: &str) {
if !self.completed_steps.contains(&step.to_string()) {
self.completed_steps.push(step.to_string());
}
self.last_updated = chrono::Local::now().format("%Y-%m-%d %H:%M:%S").to_string();
}
pub fn is_completed(&self, step: &str) -> bool {
self.completed_steps.contains(&step.to_string())
}
}
pub struct FlankingDbBuilder {
output_dir: PathBuf,
threads: usize,
amr_db_path: PathBuf,
email: String,
config: FlankBuildConfig,
}
impl FlankingDbBuilder {
pub fn new(amr_db: &Path, output_dir: &Path, threads: usize, email: &str, config: FlankBuildConfig) -> Self {
Self {
output_dir: output_dir.to_path_buf(),
threads,
amr_db_path: amr_db.to_path_buf(),
email: email.to_string(),
config,
}
}
pub fn build(&self) -> Result<()> {
eprintln!("\n============================================================");
eprintln!(" ARGenus Flanking Database Builder (Rust)");
eprintln!("============================================================");
eprintln!("Output directory: {}", self.output_dir.display());
eprintln!("NCBI Email: {}", self.email);
eprintln!("Download method: NCBI Datasets API (batch)");
eprintln!("Threads: {}", self.threads);
eprintln!("Flanking length: {} bp", self.config.flanking_length);
eprintln!("Queue buffer: {} GB", self.config.queue_buffer_gb);
eprintln!("Alignment: minimap2 (AMR indexed)");
eprintln!();
std::fs::create_dir_all(&self.output_dir)?;
let genomes_dir = self.output_dir.join("genomes");
std::fs::create_dir_all(&genomes_dir)?;
let plsdb_dir = self.config.plsdb.dir.clone().unwrap_or_else(|| self.output_dir.join("plsdb"));
if !self.config.plsdb.skip {
std::fs::create_dir_all(&plsdb_dir)?;
}
let taxonomy_dir = self.output_dir.join("taxonomy");
let temp_dir = self.output_dir.join("temp");
std::fs::create_dir_all(&temp_dir)?;
let state_path = self.output_dir.join("build_state.json");
let mut state = BuildState::load(&state_path).unwrap_or_else(|| {
eprintln!("Starting fresh build...");
BuildState::new()
});
if !state.completed_steps.is_empty() {
eprintln!("Resuming from previous state:");
eprintln!(" Started: {}", state.started_at);
eprintln!(" Last updated: {}", state.last_updated);
eprintln!(" Completed steps: {}", state.completed_steps.join(", "));
eprintln!(" Genomes downloaded: {}", state.genomes_downloaded);
eprintln!(" PLSDB extracted: {}", state.plsdb_extracted);
eprintln!();
}
let mut catalog = GenomeCatalog::new();
let catalog_path = self.output_dir.join("genome_catalog.tsv");
if !state.is_completed("taxonomy_download") {
state.update_step("taxonomy_download");
state.save(&state_path)?;
eprintln!("[1/9] Downloading NCBI taxonomy database...");
self.download_taxdump(&taxonomy_dir)?;
state.complete_step("taxonomy_download");
state.save(&state_path)?;
} else {
eprintln!("[1/9] Taxonomy database already downloaded, skipping...");
}
eprintln!("[2/9] Loading taxonomy database...");
let taxonomy = match TaxonomyDB::load(&taxonomy_dir) {
Ok(db) => {
eprintln!(" Taxonomy database loaded successfully");
Some(db)
}
Err(e) => {
eprintln!(" Warning: Failed to load taxonomy: {}", e);
eprintln!(" Genus will be parsed from organism names instead");
None
}
};
eprintln!("[3/9] Downloading NCBI assembly summaries (GenBank)...");
let assemblies = self.download_assembly_summaries()?;
eprintln!(" GenBank assemblies: {} (Complete Genome + Chromosome)", assemblies.len());
for asm in &assemblies {
catalog.add_assembly_with_taxonomy(asm, "genbank", taxonomy.as_ref());
}
let standalone_plasmids = if self.config.plsdb.skip {
eprintln!("[4/9] PLSDB skipped (--skip-plsdb)");
eprintln!("[5/9] PLSDB metadata skipped");
Vec::new()
} else if self.config.plsdb.dir.is_some() {
eprintln!("[4/9] Using pre-downloaded PLSDB: {}", plsdb_dir.display());
let nuccore_csv = plsdb_dir.join("nuccore.csv");
let fasta_path = plsdb_dir.join("sequences.fasta");
if !nuccore_csv.exists() {
let meta_tar = plsdb_dir.join("meta.tar.gz");
if meta_tar.exists() {
eprintln!(" Extracting meta.tar.gz...");
std::process::Command::new("tar")
.args(["-xzf", meta_tar.to_str().unwrap(), "-C", plsdb_dir.to_str().unwrap()])
.status()
.with_context(|| "Failed to extract PLSDB metadata")?;
} else {
anyhow::bail!("PLSDB directory missing nuccore.csv and meta.tar.gz: {}", plsdb_dir.display());
}
}
if !fasta_path.exists() {
anyhow::bail!("PLSDB directory missing sequences.fasta: {}", plsdb_dir.display());
}
state.complete_step("plsdb_download");
state.save(&state_path)?;
eprintln!("[5/9] Loading PLSDB metadata...");
let plasmids = self.load_plsdb_plasmids(&plsdb_dir)?;
eprintln!(" Standalone circular+complete plasmids: {} (not in any assembly)",
plasmids.len());
plasmids
} else {
if !state.is_completed("plsdb_download") {
state.update_step("plsdb_download");
state.save(&state_path)?;
eprintln!("[4/9] Downloading PLSDB database...");
self.download_plsdb(&plsdb_dir)?;
state.complete_step("plsdb_download");
state.save(&state_path)?;
} else {
eprintln!("[4/9] PLSDB database already downloaded, skipping...");
}
eprintln!("[5/9] Loading PLSDB metadata...");
let plasmids = self.load_plsdb_plasmids(&plsdb_dir)?;
eprintln!(" Standalone circular+complete plasmids: {} (not in any assembly)",
plasmids.len());
plasmids
};
for plasmid in &standalone_plasmids {
catalog.add_plasmid_with_taxonomy(plasmid, taxonomy.as_ref());
}
eprintln!("[6/9] Saving unified genome catalog...");
catalog.save(&catalog_path)?;
eprintln!(" Catalog entries: {}", catalog.len());
eprintln!(" Saved to: {}", catalog_path.display());
let paf_output = self.output_dir.join("all_alignments.paf");
let merged_hits = self.output_dir.join("merged_alignment_hits.tsv");
let output_tsv = self.output_dir.join("all_flanking_sequences.tsv");
if !state.is_completed("alignment_flanking") {
state.update_step("genome_download");
state.save(&state_path)?;
eprintln!("[7/9] Batch processing: Download + Align NCBI genomes...");
eprintln!(" Producer-consumer pattern with {} GB queue buffer", self.config.queue_buffer_gb);
let downloaded = self.download_and_align_batches(
&assemblies,
&genomes_dir,
&temp_dir,
&paf_output,
&catalog,
)?;
state.genomes_downloaded = downloaded;
state.complete_step("genome_download");
state.save(&state_path)?;
eprintln!(" Processed {} genomes", downloaded);
if self.config.plsdb.skip {
eprintln!("[8/9] PLSDB processing skipped (--skip-plsdb)");
state.plsdb_extracted = 0;
} else {
state.update_step("plsdb_extract");
state.save(&state_path)?;
eprintln!("[8/9] Processing PLSDB sequences...");
let plsdb_count = self.align_plsdb(
&standalone_plasmids,
&plsdb_dir,
&temp_dir,
&paf_output,
)?;
state.plsdb_extracted = plsdb_count;
state.complete_step("plsdb_extract");
state.save(&state_path)?;
}
state.update_step("alignment_flanking");
state.save(&state_path)?;
eprintln!("[9/9] Processing alignment results...");
self.convert_paf_to_merged(&paf_output, &merged_hits)?;
state.alignments_done = true;
state.save(&state_path)?;
eprintln!(" Extracting flanking sequences (flanking_length: {} bp)...",
self.config.flanking_length);
self.extract_flanking_sequences(
&merged_hits,
&genomes_dir,
&output_tsv,
&catalog,
)?;
if !self.config.plsdb.skip {
self.extract_flanking_from_plsdb(
&merged_hits,
&plsdb_dir,
&output_tsv,
&catalog,
)?;
}
state.flanking_extracted = true;
state.complete_step("alignment_flanking");
state.save(&state_path)?;
eprintln!(" Cleaning up PAF file...");
std::fs::remove_file(&paf_output).ok();
} else {
eprintln!("[7-9] Batch pipeline already completed, skipping...");
}
if temp_dir.exists() {
eprintln!("\nCleaning up temp files...");
std::fs::remove_dir_all(&temp_dir).ok();
}
let fdb_path = self.output_dir.join("flanking.fdb");
if !state.is_completed("fdb_build") && output_tsv.exists() {
state.update_step("fdb_build");
state.save(&state_path)?;
eprintln!("\n[10/10] Building FDB from TSV (external sort + zstd)...");
let buffer_size_mb = 1024;
crate::fdb::build(&output_tsv, &fdb_path, buffer_size_mb, self.threads)?;
state.complete_step("fdb_build");
state.save(&state_path)?;
} else if fdb_path.exists() {
eprintln!("[10/10] FDB already built, skipping...");
}
eprintln!("\n============================================================");
eprintln!(" Build Complete");
eprintln!("============================================================");
eprintln!("Output files:");
eprintln!(" - {}", output_tsv.display());
if fdb_path.exists() {
let fdb_size = std::fs::metadata(&fdb_path).map(|m| m.len()).unwrap_or(0);
eprintln!(" - {} ({:.1} MB)", fdb_path.display(), fdb_size as f64 / 1024.0 / 1024.0);
}
eprintln!(" - {}", catalog_path.display());
Ok(())
}
fn download_assembly_summaries(&self) -> Result<Vec<AssemblyInfo>> {
let mut assemblies = Vec::new();
for kingdom in ["bacteria", "archaea"] {
let url = format!(
"{}/genbank/{}/assembly_summary.txt",
NCBI_FTP_BASE, kingdom
);
eprintln!(" Downloading genbank/{}...", kingdom);
match self.download_and_parse_assembly_summary(&url, None) {
Ok((mut asm, _)) => {
eprintln!(" Found {} assemblies (Complete Genome + Chromosome)", asm.len());
assemblies.append(&mut asm);
}
Err(e) => {
eprintln!(" Warning: Failed to download {}: {}", url, e);
}
}
}
eprintln!(" Total assemblies: {}", assemblies.len());
Ok(assemblies)
}
fn download_plsdb(&self, plsdb_dir: &Path) -> Result<()> {
let meta_tar = plsdb_dir.join("meta.tar.gz");
let fasta_path = plsdb_dir.join("sequences.fasta");
if !plsdb_dir.join("nuccore.csv").exists() {
eprintln!(" Downloading PLSDB metadata...");
self.download_file(PLSDB_META_URL, &meta_tar)?;
eprintln!(" Extracting metadata...");
let status = Command::new("tar")
.args(["-xzf", meta_tar.to_str().unwrap(), "-C", plsdb_dir.to_str().unwrap()])
.status()
.with_context(|| "Failed to extract PLSDB metadata")?;
if !status.success() {
anyhow::bail!("tar extraction failed");
}
std::fs::remove_file(&meta_tar).ok();
} else {
eprintln!(" PLSDB metadata already exists, skipping download...");
}
if !fasta_path.exists() {
eprintln!(" Downloading PLSDB sequences (~7GB)...");
self.download_file(PLSDB_FASTA_URL, &fasta_path)?;
} else {
eprintln!(" PLSDB sequences already exist, skipping download...");
}
Ok(())
}
fn download_taxdump(&self, taxonomy_dir: &Path) -> Result<()> {
let tar_path = taxonomy_dir.join("taxdump.tar.gz");
let names_path = taxonomy_dir.join("names.dmp");
if names_path.exists() {
eprintln!(" Taxdump already exists, skipping download...");
return Ok(());
}
std::fs::create_dir_all(taxonomy_dir)?;
eprintln!(" Downloading NCBI taxdump (~60MB)...");
self.download_file(NCBI_TAXDUMP_URL, &tar_path)?;
eprintln!(" Extracting taxdump...");
let status = Command::new("tar")
.args(["-xzf", tar_path.to_str().unwrap(), "-C", taxonomy_dir.to_str().unwrap()])
.status()
.with_context(|| "Failed to extract taxdump")?;
if !status.success() {
anyhow::bail!("tar extraction failed");
}
std::fs::remove_file(&tar_path).ok();
Ok(())
}
fn download_file(&self, url: &str, output_path: &Path) -> Result<()> {
for attempt in 0..3 {
match self.download_file_once(url, output_path) {
Ok(_) => return Ok(()),
Err(e) if attempt < 2 => {
eprintln!(" Download failed (attempt {}): {}", attempt + 1, e);
eprintln!(" Retrying in 5 seconds...");
std::thread::sleep(Duration::from_secs(5));
continue;
}
Err(e) => return Err(e),
}
}
Ok(())
}
fn download_file_once(&self, url: &str, output_path: &Path) -> Result<()> {
let response = ureq::get(url)
.set("User-Agent", "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36")
.timeout(Duration::from_secs(7200)) .call()
.with_context(|| format!("Failed to download {}", url))?;
let mut file = File::create(output_path)?;
let mut reader = response.into_reader();
let mut buffer = [0u8; 65536];
let mut total = 0usize;
loop {
match reader.read(&mut buffer) {
Ok(0) => break,
Ok(n) => {
file.write_all(&buffer[..n])?;
total += n;
if total.is_multiple_of(100 * 1024 * 1024) {
eprintln!(" Downloaded {} MB...", total / (1024 * 1024));
}
}
Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
Err(e) => return Err(e.into()),
}
}
eprintln!(" Total: {} MB", total / (1024 * 1024));
Ok(())
}
fn download_and_parse_assembly_summary(
&self,
url: &str,
exclude_set: Option<&FxHashSet<String>>,
) -> Result<(Vec<AssemblyInfo>, Vec<String>)> {
let response = ureq::get(url)
.timeout(Duration::from_secs(600))
.call()
.with_context(|| format!("Failed to download {}", url))?;
let reader = BufReader::new(response.into_reader());
let mut assemblies = Vec::new();
let mut paired_gca = Vec::new();
let mut skipped_duplicates = 0usize;
for line in reader.lines() {
let line = line?;
if line.starts_with('#') {
continue;
}
let fields: Vec<&str> = line.split('\t').collect();
if fields.len() < 20 {
continue;
}
let accession = fields[0];
let assembly_level = fields[11];
if assembly_level != "Complete Genome" && assembly_level != "Chromosome" {
continue;
}
let ftp_path = fields[19];
if ftp_path == "na" || ftp_path.is_empty() {
continue;
}
if let Some(exclude) = exclude_set {
let acc_base = accession.split('.').next().unwrap_or(accession);
if exclude.contains(accession) || exclude.contains(acc_base) {
skipped_duplicates += 1;
continue;
}
}
if fields.len() > 18 && accession.starts_with("GCF_") {
let paired_asm = fields[17];
let paired_comp = fields[18];
if paired_asm != "na" && !paired_asm.is_empty() && paired_comp == "identical" {
paired_gca.push(paired_asm.to_string());
if let Some(base) = paired_asm.split('.').next() {
paired_gca.push(base.to_string());
}
}
}
assemblies.push(AssemblyInfo {
accession: accession.to_string(),
taxid: fields[5].to_string(),
species_taxid: fields[6].to_string(),
organism_name: fields[7].to_string(),
});
}
if skipped_duplicates > 0 {
eprintln!(" Skipped {} duplicates (have identical RefSeq)", skipped_duplicates);
}
Ok((assemblies, paired_gca))
}
fn load_plsdb_plasmids(&self, plsdb_dir: &Path) -> Result<Vec<PlasmidInfo>> {
let nuccore_path = plsdb_dir.join("nuccore.csv");
let taxonomy_path = plsdb_dir.join("taxonomy.csv");
let mut taxonomy: FxHashMap<String, (String, String)> = FxHashMap::default();
if taxonomy_path.exists() {
let tax_file = File::open(&taxonomy_path)?;
let tax_reader = BufReader::new(tax_file);
for (idx, line) in tax_reader.lines().enumerate() {
let line = line?;
if idx == 0 {
continue; }
let fields: Vec<&str> = line.split(',').collect();
if fields.len() >= 10 {
let uid = fields[0].to_string();
let genus = fields[8].to_string();
let species = fields[9].to_string();
taxonomy.insert(uid, (genus, species));
}
}
}
let mut plasmids = Vec::new();
let nuc_file = File::open(&nuccore_path)?;
let nuc_reader = BufReader::new(nuc_file);
for (idx, line) in nuc_reader.lines().enumerate() {
let line = line?;
if idx == 0 {
continue; }
let fields = parse_csv_line(&line);
if fields.len() < 15 {
continue;
}
let completeness = &fields[4];
let topology = &fields[14];
let assembly_uid = &fields[9];
if completeness != "complete" {
continue;
}
if topology != "circular" {
continue;
}
if assembly_uid != "-1" {
continue;
}
let taxonomy_uid = &fields[12];
let (genus, species) = taxonomy.get(taxonomy_uid)
.cloned()
.unwrap_or_else(|| ("Unknown".to_string(), "Unknown".to_string()));
plasmids.push(PlasmidInfo {
accession: fields[1].clone(),
taxonomy_uid: taxonomy_uid.clone(),
genus,
species,
});
}
Ok(plasmids)
}
fn download_and_align_batches(
&self,
assemblies: &[AssemblyInfo],
genomes_dir: &Path,
temp_dir: &Path,
paf_output: &Path,
_catalog: &GenomeCatalog,
) -> Result<usize> {
let all_accessions_path = temp_dir.join("accessions_all.txt");
let done_accessions_path = temp_dir.join("accessions_done.txt");
let accessions: Vec<String> = assemblies.iter()
.map(|a| a.accession.clone())
.collect();
if !all_accessions_path.exists() {
eprintln!(" Saving {} accessions to temp file...", accessions.len());
let mut file = File::create(&all_accessions_path)?;
for acc in &accessions {
writeln!(file, "{}", acc)?;
}
}
let mut done_set: FxHashSet<String> = FxHashSet::default();
if done_accessions_path.exists() {
let file = File::open(&done_accessions_path)?;
let reader = BufReader::new(file);
for acc in reader.lines().map_while(Result::ok) {
done_set.insert(acc.trim().to_string());
}
if !done_set.is_empty() {
if !paf_output.exists() {
eprintln!(" WARNING: PAF file missing but {} genomes marked as done", done_set.len());
eprintln!(" Clearing done list and reprocessing all genomes...");
done_set.clear();
std::fs::remove_file(&done_accessions_path).ok();
} else {
let paf_size = std::fs::metadata(paf_output)?.len();
if paf_size == 0 {
eprintln!(" WARNING: PAF file is empty but {} genomes marked as done", done_set.len());
eprintln!(" Clearing done list and reprocessing all genomes...");
done_set.clear();
std::fs::remove_file(&done_accessions_path).ok();
} else {
let mut missing_genomes = Vec::new();
for acc in &done_set {
let genome_path = genomes_dir.join(format!("{}.fna", acc));
if !genome_path.exists() {
missing_genomes.push(acc.clone());
}
}
if !missing_genomes.is_empty() {
eprintln!(" WARNING: {} genome files missing for done accessions", missing_genomes.len());
for acc in &missing_genomes {
done_set.remove(acc);
}
eprintln!(" Removed missing genomes from done list, will re-download");
}
if !done_set.is_empty() {
eprintln!(" Resume: {} already processed (PAF: {} MB)",
done_set.len(), paf_size / 1024 / 1024);
}
}
}
}
}
let remaining: Vec<&String> = accessions.iter()
.filter(|a| !done_set.contains(*a))
.collect();
if remaining.is_empty() {
eprintln!(" All {} genomes already processed", accessions.len());
return Ok(accessions.len());
}
eprintln!(" Remaining: {}/{} genomes to process", remaining.len(), accessions.len());
let bytes_per_batch = API_BATCH_SIZE * 3 * 1024 * 1024; let queue_capacity = std::cmp::max(
1,
(self.config.queue_buffer_gb as usize * 1024 * 1024 * 1024) / bytes_per_batch
);
eprintln!(" Queue capacity: {} batches (based on {} GB buffer)", queue_capacity, self.config.queue_buffer_gb);
let batches: Vec<Vec<String>> = remaining.chunks(API_BATCH_SIZE)
.map(|chunk| chunk.iter().map(|s| (*s).clone()).collect())
.collect();
let total_batches = batches.len();
let (tx, rx) = mpsc::sync_channel::<(usize, Vec<PathBuf>)>(queue_capacity);
let downloaded = Arc::new(AtomicUsize::new(done_set.len()));
let aligned = Arc::new(AtomicUsize::new(done_set.len()));
let download_error = Arc::new(AtomicBool::new(false));
let genomes_dir = genomes_dir.to_path_buf();
let temp_dir_producer = temp_dir.to_path_buf();
let temp_dir_consumer = temp_dir.to_path_buf();
let email = self.email.clone();
let downloaded_clone = Arc::clone(&downloaded);
let download_error_clone = Arc::clone(&download_error);
let total_accessions = accessions.len();
let producer = thread::spawn(move || -> Result<()> {
for (batch_idx, batch) in batches.into_iter().enumerate() {
if download_error_clone.load(Ordering::Relaxed) {
break;
}
let zip_path = temp_dir_producer.join(format!("batch_{:04}.zip", batch_idx));
let url = format!("{}/genome/download", NCBI_DATASETS_API);
let acc_list: Vec<&str> = batch.iter().map(|s| s.as_str()).collect();
let request_body = serde_json::json!({
"accessions": acc_list,
"include_annotation_type": ["GENOME_FASTA"]
});
let response = ureq::post(&url)
.set("Content-Type", "application/json")
.set("ncbi-client-id", &email)
.timeout(Duration::from_secs(3600))
.send_json(&request_body);
match response {
Ok(resp) => {
let mut zip_file = File::create(&zip_path)?;
let mut reader = resp.into_reader();
std::io::copy(&mut reader, &mut zip_file)?;
drop(zip_file);
let genome_files = Self::extract_batch_to_files_static(&zip_path, &genomes_dir)?;
std::fs::remove_file(&zip_path).ok();
downloaded_clone.fetch_add(batch.len(), Ordering::Relaxed);
eprintln!(" [Download] Batch {}/{}: {} genomes ({}/{})",
batch_idx + 1, total_batches, batch.len(),
downloaded_clone.load(Ordering::Relaxed), total_accessions);
if tx.send((batch_idx, genome_files)).is_err() {
break; }
}
Err(e) => {
eprintln!(" [Download] Batch {} failed: {}", batch_idx + 1, e);
download_error_clone.store(true, Ordering::Relaxed);
break;
}
}
std::thread::sleep(Duration::from_millis(200));
}
Ok(())
});
let paf_file = Arc::new(Mutex::new(
std::fs::OpenOptions::new()
.create(true)
.append(true)
.open(paf_output)?
));
let done_file = Arc::new(Mutex::new(
std::fs::OpenOptions::new()
.create(true)
.append(true)
.open(&done_accessions_path)?
));
let amr_db = self.amr_db_path.clone();
let max_threads = self.threads;
while let Ok((batch_idx, genome_files)) = rx.recv() {
if genome_files.is_empty() {
continue;
}
let buckets = bucket_genomes_by_size(&genome_files);
let mut batch_hits: Vec<PafHit> = Vec::new();
for (category, bucket_files) in &buckets {
if bucket_files.is_empty() {
continue;
}
let bucket_threads = category.thread_count(max_threads);
let bucket_suffix = format!("batch_{:04}_{:?}", batch_idx, category);
let bucket_fasta = temp_dir_consumer.join(format!("{}.fas", bucket_suffix));
{
let mut fasta_writer = BufWriter::new(File::create(&bucket_fasta)?);
for genome_path in bucket_files {
let filename = genome_path.file_name()
.and_then(|n| n.to_str())
.unwrap_or("unknown.fna");
let file = File::open(genome_path)?;
let reader = BufReader::new(file);
for line in reader.lines() {
let line = line?;
if let Some(stripped) = line.strip_prefix('>') {
let contig_id = stripped.split_whitespace().next().unwrap_or(stripped);
writeln!(fasta_writer, ">{}|{}", contig_id, filename)?;
} else {
writeln!(fasta_writer, "{}", line)?;
}
}
}
fasta_writer.flush()?;
}
let bucket_paf = temp_dir_consumer.join(format!("{}.paf", bucket_suffix));
let status = Command::new("minimap2")
.args([
"-cx", "asm20",
"-t", &bucket_threads.to_string(),
amr_db.to_str().unwrap(),
bucket_fasta.to_str().unwrap(),
"-o", bucket_paf.to_str().unwrap(),
])
.stdout(std::process::Stdio::null())
.stderr(std::process::Stdio::null())
.status();
if let Ok(s) = status {
if s.success() && bucket_paf.exists() {
let paf_content = std::fs::read_to_string(&bucket_paf)?;
let hits: Vec<PafHit> = paf_content
.lines()
.filter_map(PafHit::from_paf_line)
.collect();
batch_hits.extend(hits);
}
}
std::fs::remove_file(&bucket_fasta).ok();
std::fs::remove_file(&bucket_paf).ok();
}
let dedup_hits = deduplicate_paf_hits(batch_hits);
{
let mut paf = paf_file.lock().unwrap();
for hit in &dedup_hits {
writeln!(paf, "{}", hit.raw_line)?;
}
paf.flush()?; }
{
let mut done = done_file.lock().unwrap();
for genome_path in &genome_files {
if let Some(stem) = genome_path.file_stem().and_then(|s| s.to_str()) {
writeln!(done, "{}", stem)?;
}
}
done.flush()?; }
let count = genome_files.len();
aligned.fetch_add(count, Ordering::Relaxed);
let bucket_info: Vec<String> = buckets.iter()
.map(|(cat, files)| format!("{}={}", cat.name().split_whitespace().next().unwrap_or("?"), files.len()))
.collect();
eprintln!(" [Align] Batch {}/{}: {} genomes [{}] ({}/{})",
batch_idx + 1, total_batches, count,
bucket_info.join(", "),
aligned.load(Ordering::Relaxed), total_accessions);
}
let _ = producer.join();
Ok(aligned.load(Ordering::Relaxed))
}
fn extract_batch_to_files_static(zip_path: &Path, genomes_dir: &Path) -> Result<Vec<PathBuf>> {
let zip_file = File::open(zip_path)?;
let mut archive = zip::ZipArchive::new(zip_file)?;
let mut genome_files = Vec::new();
for i in 0..archive.len() {
let mut file = archive.by_index(i)?;
let name = file.name().to_string();
if name.ends_with("_genomic.fna") || name.ends_with("_genomic.fasta") {
let parts: Vec<&str> = name.split('/').collect();
if let Some(acc_dir) = parts.iter().find(|p| p.starts_with("GCA_") || p.starts_with("GCF_")) {
let filename = format!("{}.fna", acc_dir);
let output_path = genomes_dir.join(&filename);
let mut content = Vec::new();
file.read_to_end(&mut content)?;
std::fs::write(&output_path, &content)?;
genome_files.push(output_path);
}
}
}
Ok(genome_files)
}
fn align_plsdb(
&self,
plasmids: &[PlasmidInfo],
plsdb_dir: &Path,
temp_dir: &Path,
paf_output: &Path,
) -> Result<usize> {
let fasta_path = plsdb_dir.join("sequences.fasta");
if !fasta_path.exists() {
eprintln!(" Warning: sequences.fasta not found in PLSDB directory");
return Ok(0);
}
let target_accs: FxHashSet<_> = plasmids.iter()
.map(|p| p.accession.clone())
.collect();
if target_accs.is_empty() {
eprintln!(" No PLSDB plasmids to process");
return Ok(0);
}
let acc_list: Vec<_> = target_accs.iter().collect();
let batches: Vec<_> = acc_list.chunks(API_BATCH_SIZE).collect();
let total_batches = batches.len();
let mut total_processed = 0usize;
eprintln!(" Processing {} PLSDB plasmids in {} batches...",
target_accs.len(), total_batches);
let file = File::open(&fasta_path)?;
let reader = BufReader::new(file);
let mut seq_map: FxHashMap<String, String> = FxHashMap::default();
let mut current_acc: Option<String> = None;
let mut current_seq = String::new();
for line in reader.lines() {
let line = line?;
if line.starts_with('>') {
if let Some(acc) = current_acc.take() {
if target_accs.contains(&acc) {
seq_map.insert(acc, std::mem::take(&mut current_seq));
}
}
let header = line.trim_start_matches('>');
let acc = header.split_whitespace().next().unwrap_or("").to_string();
current_acc = Some(acc);
current_seq.clear();
} else {
current_seq.push_str(line.trim());
}
}
if let Some(acc) = current_acc {
if target_accs.contains(&acc) {
seq_map.insert(acc, current_seq);
}
}
eprintln!(" Loaded {} target sequences from PLSDB", seq_map.len());
for (batch_idx, batch) in batches.iter().enumerate() {
let mut size_buckets: FxHashMap<GenomeSizeCategory, Vec<(&str, &str)>> = FxHashMap::default();
for acc in *batch {
if let Some(seq) = seq_map.get(*acc) {
let category = GenomeSizeCategory::from_size(seq.len() as u64);
size_buckets.entry(category)
.or_default()
.push((acc.as_str(), seq.as_str()));
}
}
let mut batch_hits: Vec<PafHit> = Vec::new();
for (category, bucket_seqs) in &size_buckets {
if bucket_seqs.is_empty() {
continue;
}
let bucket_threads = category.thread_count(self.threads);
let bucket_suffix = format!("plsdb_batch_{:04}_{:?}", batch_idx, category);
let bucket_fasta = temp_dir.join(format!("{}.fas", bucket_suffix));
{
let mut fasta_writer = BufWriter::new(File::create(&bucket_fasta)?);
for (acc, seq) in bucket_seqs {
let filename = format!("{}.fna", acc.replace('.', "_"));
writeln!(fasta_writer, ">{}|{}", acc, filename)?;
writeln!(fasta_writer, "{}", seq)?;
}
fasta_writer.flush()?;
}
let bucket_paf = temp_dir.join(format!("{}.paf", bucket_suffix));
let status = Command::new("minimap2")
.args([
"-cx", "asm20",
"-t", &bucket_threads.to_string(),
self.amr_db_path.to_str().unwrap(),
bucket_fasta.to_str().unwrap(),
"-o", bucket_paf.to_str().unwrap(),
])
.stdout(std::process::Stdio::null())
.stderr(std::process::Stdio::null())
.status();
if let Ok(s) = status {
if s.success() && bucket_paf.exists() {
let paf_content = std::fs::read_to_string(&bucket_paf)?;
let hits: Vec<PafHit> = paf_content
.lines()
.filter_map(PafHit::from_paf_line)
.collect();
batch_hits.extend(hits);
}
}
std::fs::remove_file(&bucket_fasta).ok();
std::fs::remove_file(&bucket_paf).ok();
}
let dedup_hits = deduplicate_paf_hits(batch_hits);
{
let mut paf_file = std::fs::OpenOptions::new()
.create(true)
.append(true)
.open(paf_output)?;
for hit in &dedup_hits {
writeln!(paf_file, "{}", hit.raw_line)?;
}
}
total_processed += batch.len();
let bucket_info: Vec<String> = size_buckets.iter()
.map(|(cat, seqs)| format!("{}={}", cat.name().split_whitespace().next().unwrap_or("?"), seqs.len()))
.collect();
eprintln!(" [PLSDB] Batch {}/{}: {} plasmids [{}] ({}/{})",
batch_idx + 1, total_batches, batch.len(),
bucket_info.join(", "),
total_processed, target_accs.len());
}
for (acc, seq) in &seq_map {
let out_path = plsdb_dir.join(format!("{}.fna", acc.replace('.', "_")));
if !out_path.exists() {
let mut out_file = File::create(&out_path)?;
writeln!(out_file, ">{}", acc)?;
writeln!(out_file, "{}", seq)?;
}
}
eprintln!(" Aligned {} PLSDB sequences", total_processed);
Ok(total_processed)
}
fn extract_flanking_from_plsdb(
&self,
hits_path: &Path,
plsdb_dir: &Path,
output_path: &Path,
catalog: &GenomeCatalog,
) -> Result<()> {
let hits_file = File::open(hits_path)?;
let reader = BufReader::new(hits_file);
let mut plsdb_hits: FxHashMap<String, Vec<(String, String, usize, usize)>> = FxHashMap::default();
for (i, line) in reader.lines().enumerate() {
let line = line?;
if i == 0 && line.starts_with("gene") {
continue;
}
let fields: Vec<&str> = line.split('\t').collect();
if fields.len() < 4 {
continue;
}
let gene = fields[0];
let contig_file = fields[1];
let start: usize = fields[2].parse().unwrap_or(0);
let end: usize = fields[3].parse().unwrap_or(0);
let (contig_id, genome_file) = if let Some(pipe_pos) = contig_file.rfind('|') {
(contig_file[..pipe_pos].to_string(), contig_file[pipe_pos + 1..].to_string())
} else {
continue;
};
if genome_file.starts_with("NZ_") || genome_file.starts_with("CP") ||
genome_file.starts_with("AP") || genome_file.starts_with("NC_") {
plsdb_hits.entry(genome_file)
.or_default()
.push((gene.to_string(), contig_id, start, end));
}
}
if plsdb_hits.is_empty() {
return Ok(());
}
let mut writer = std::fs::OpenOptions::new()
.create(true)
.append(true)
.open(output_path)?;
let mut extracted = 0usize;
for (genome_file, hits) in &plsdb_hits {
let genome_path = plsdb_dir.join(genome_file);
if !genome_path.exists() {
continue;
}
let sequences = match self.load_genome_sequences(&genome_path) {
Ok(seqs) => seqs,
Err(_) => continue,
};
let seq_map: FxHashMap<&str, &str> = sequences.iter()
.map(|(header, seq)| {
let key = header.split_whitespace().next().unwrap_or(header.as_str());
(key, seq.as_str())
})
.collect();
let genome_acc = genome_file.trim_end_matches(".fna");
let base_genus = catalog.get_genus(genome_acc).map(|s| s.to_string());
for (gene, contig_id, start, end) in hits {
let contig_seq = match seq_map.get(contig_id.as_str()) {
Some(seq) => *seq,
None => continue,
};
let contig_len = contig_seq.len();
if *start >= contig_len || *end > contig_len || start >= end {
continue;
}
let genus = base_genus.clone().unwrap_or_else(|| "Unknown".to_string());
let upstream_start = start.saturating_sub(self.config.flanking_length);
let upstream = &contig_seq[upstream_start..*start];
let downstream_end = std::cmp::min(*end + self.config.flanking_length, contig_len);
let downstream = &contig_seq[*end..downstream_end];
writeln!(writer, "{}\t{}\t{}\t{}\t{}\t{}\t{}",
gene, contig_id, genus, start, end, upstream, downstream)?;
extracted += 1;
}
}
if extracted > 0 {
eprintln!(" Extracted {} PLSDB flanking sequences", extracted);
}
Ok(())
}
fn convert_paf_to_merged(
&self,
paf_path: &Path,
output_path: &Path,
) -> Result<()> {
let mut hits: FxHashSet<(String, String, usize, usize)> = FxHashSet::default();
if paf_path.exists() {
let file = File::open(paf_path)?;
let reader = BufReader::new(file);
for line in reader.lines() {
let line = line?;
let fields: Vec<&str> = line.split('\t').collect();
if fields.len() < 12 {
continue;
}
let contig_file = fields[0].to_string(); let gene = fields[5].to_string(); let start: usize = fields[2].parse().unwrap_or(0);
let end: usize = fields[3].parse().unwrap_or(0);
hits.insert((gene, contig_file, start, end));
}
}
let mut writer = BufWriter::new(File::create(output_path)?);
for (gene, contig_file, start, end) in &hits {
writeln!(writer, "{}\t{}\t{}\t{}", gene, contig_file, start, end)?;
}
eprintln!(" Found {} unique hits", hits.len());
Ok(())
}
fn extract_flanking_sequences(
&self,
hits_path: &Path,
genomes_dir: &Path,
output_path: &Path,
catalog: &GenomeCatalog,
) -> Result<()> {
if output_path.exists() {
let metadata = std::fs::metadata(output_path)?;
if metadata.len() > 0 {
eprintln!(" Flanking sequences file already exists ({} MB), skipping extraction...",
metadata.len() / (1024 * 1024));
return Ok(());
}
}
let hits_file = File::open(hits_path)?;
let reader = BufReader::new(hits_file);
let mut genome_hits: FxHashMap<String, Vec<(String, String, usize, usize)>> = FxHashMap::default();
for (i, line) in reader.lines().enumerate() {
let line = line?;
if i == 0 && line.starts_with("gene") {
continue;
}
let fields: Vec<&str> = line.split('\t').collect();
if fields.len() < 4 {
continue;
}
let gene = fields[0];
let contig_file = fields[1]; let start: usize = fields[2].parse().unwrap_or(0);
let end: usize = fields[3].parse().unwrap_or(0);
let (contig_id, genome_file) = if let Some(pipe_pos) = contig_file.rfind('|') {
(contig_file[..pipe_pos].to_string(), contig_file[pipe_pos + 1..].to_string())
} else {
(contig_file.to_string(), extract_genome_file(contig_file))
};
genome_hits.entry(genome_file)
.or_default()
.push((gene.to_string(), contig_id, start, end));
}
let output_file = File::create(output_path)?;
let mut writer = BufWriter::new(output_file);
writeln!(writer, "Gene\tContig\tGenus\tStart\tEnd\tUpstream\tDownstream")?;
let total_genomes = genome_hits.len();
let mut processed = 0;
let mut genus_found = 0usize;
let mut genus_unknown = 0usize;
let mut sequences_extracted = 0usize;
for (genome_file, hits) in &genome_hits {
let genome_path = genomes_dir.join(genome_file);
if !genome_path.exists() {
continue;
}
let sequences = match self.load_genome_sequences(&genome_path) {
Ok(seqs) => seqs,
Err(_) => continue,
};
let seq_map: FxHashMap<&str, &str> = sequences.iter()
.map(|(header, seq)| {
let key = header.split_whitespace().next().unwrap_or(header.as_str());
(key, seq.as_str())
})
.collect();
let genome_acc = genome_file.trim_end_matches(".fna");
let base_genus = catalog.get_genus(genome_acc).map(|s| s.to_string());
for (gene, contig_id, start, end) in hits {
let contig_seq = match seq_map.get(contig_id.as_str()) {
Some(seq) => *seq,
None => continue,
};
let contig_len = contig_seq.len();
if *start >= contig_len || *end > contig_len || start >= end {
continue;
}
let upstream_start = start.saturating_sub(self.config.flanking_length);
let downstream_end = std::cmp::min(end + self.config.flanking_length, contig_len);
let upstream = if *start > 0 && upstream_start < *start {
&contig_seq[upstream_start..*start]
} else {
""
};
let downstream = if *end < contig_len && *end < downstream_end {
&contig_seq[*end..downstream_end]
} else {
""
};
let genus = if let Some(ref g) = base_genus {
genus_found += 1;
g.clone()
} else if let Some(g) = catalog.get_genus(contig_id) {
genus_found += 1;
g.to_string()
} else {
genus_unknown += 1;
"Unknown".to_string()
};
writeln!(
writer,
"{}\t{}\t{}\t{}\t{}\t{}\t{}",
gene, contig_id, genus, start, end, upstream, downstream
)?;
sequences_extracted += 1;
}
processed += 1;
if processed % 10000 == 0 || processed == total_genomes {
eprintln!(" Processed {}/{} genomes ({} sequences)",
processed, total_genomes, sequences_extracted);
}
}
eprintln!(" Extracted {} flanking sequences from {} genomes",
sequences_extracted, processed);
eprintln!(" Genus resolved: {}, Unknown: {}", genus_found, genus_unknown);
Ok(())
}
fn load_genome_sequences(&self, genome_path: &Path) -> Result<Vec<(String, String)>> {
let file = File::open(genome_path)?;
let reader = BufReader::new(file);
let mut sequences = Vec::new();
let mut current_name: Option<String> = None;
let mut current_seq = String::new();
for line in reader.lines() {
let line = line?;
if let Some(stripped) = line.strip_prefix('>') {
if let Some(name) = current_name.take() {
sequences.push((name, std::mem::take(&mut current_seq)));
}
current_name = Some(stripped.to_string());
current_seq.clear();
} else {
current_seq.push_str(line.trim());
}
}
if let Some(name) = current_name {
sequences.push((name, current_seq));
}
Ok(sequences)
}
}
fn parse_csv_line(line: &str) -> Vec<String> {
let mut fields = Vec::new();
let mut current = String::new();
let mut in_quotes = false;
for ch in line.chars() {
match ch {
'"' => in_quotes = !in_quotes,
',' if !in_quotes => {
fields.push(std::mem::take(&mut current));
}
_ => current.push(ch),
}
}
fields.push(current);
fields
}
fn extract_genome_file(contig: &str) -> String {
if let Some(start) = contig.find("GC") {
let remainder = &contig[start..];
if let Some(end) = remainder.find(|c: char| !c.is_alphanumeric() && c != '_' && c != '.') {
return format!("{}.fna", &remainder[..end]);
}
return format!("{}.fna", remainder);
}
let acc = contig.split_whitespace().next().unwrap_or(contig);
format!("{}.fna", acc.replace('.', "_"))
}
pub fn build(output_dir: &Path, arg_db: &Path, threads: usize, email: &str, config: FlankBuildConfig) -> Result<()> {
let builder = FlankingDbBuilder::new(arg_db, output_dir, threads, email, config);
builder.build()
}