src/preprocessing/amr_gene_detector.py

"""AMR gene detection using ResFinder database.

This module detects antimicrobial resistance genes in genome sequences
by searching for ResFinder reference sequences. No external tools required.
"""

import gzip
import json
import logging
from collections import defaultdict
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple

import numpy as np
import pandas as pd

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class AMRGeneDetector:
    """Detect AMR genes in genomes using ResFinder database.

    Uses k-mer based sequence matching to identify resistance genes
    without requiring external alignment tools.
    """

    def __init__(
        self,
        resfinder_dir: str = "data/raw/resfinder",
        card_dir: str = "data/raw/card-data",
        kmer_size: int = 31,
        min_identity: float = 0.8,
    ):
        """Initialize AMR gene detector.

        Args:
            resfinder_dir: Path to ResFinder database
            card_dir: Path to CARD database
            kmer_size: Size of k-mers for matching (default 31)
            min_identity: Minimum fraction of k-mers that must match (default 0.8)
        """
        self.resfinder_dir = Path(resfinder_dir)
        self.card_dir = Path(card_dir)
        self.kmer_size = kmer_size
        self.min_identity = min_identity

        # Reference data
        self.resistance_genes: Dict[str, Dict] = {}  # gene_name -> {sequence, drug_class, ...}
        self.drug_classes: Set[str] = set()
        self.gene_kmers: Dict[str, Set[str]] = {}  # gene_name -> set of k-mers

    def load_resfinder_database(self) -> None:
        """Load resistance genes from ResFinder database."""
        logger.info("Loading ResFinder database...")

        # Map of file names to drug classes
        drug_class_files = {
            "aminoglycoside.fsa": "aminoglycoside",
            "beta-lactam.fsa": "beta-lactam",
            "colistin.fsa": "colistin",
            "fosfomycin.fsa": "fosfomycin",
            "fusidicacid.fsa": "fusidic_acid",
            "glycopeptide.fsa": "glycopeptide",
            "macrolide.fsa": "macrolide",
            "nitroimidazole.fsa": "nitroimidazole",
            "oxazolidinone.fsa": "oxazolidinone",
            "phenicol.fsa": "phenicol",
            "quinolone.fsa": "quinolone",
            "rifampicin.fsa": "rifampicin",
            "sulphonamide.fsa": "sulfonamide",
            "tetracycline.fsa": "tetracycline",
            "trimethoprim.fsa": "trimethoprim",
        }

        for filename, drug_class in drug_class_files.items():
            fasta_path = self.resfinder_dir / filename
            if fasta_path.exists():
                genes = self._parse_fasta(fasta_path)
                for gene_name, sequence in genes.items():
                    self.resistance_genes[gene_name] = {
                        "sequence": sequence,
                        "drug_class": drug_class,
                        "source": "resfinder",
                    }
                    self.drug_classes.add(drug_class)
                logger.info(f"  Loaded {len(genes)} genes from {filename}")

        logger.info(f"Total resistance genes loaded: {len(self.resistance_genes)}")
        logger.info(f"Drug classes: {sorted(self.drug_classes)}")

    def _parse_fasta(self, fasta_path: Path) -> Dict[str, str]:
        """Parse a FASTA file and return gene sequences."""
        genes = {}
        current_gene = None
        current_seq = []

        with open(fasta_path, "r") as f:
            for line in f:
                line = line.strip()
                if line.startswith(">"):
                    if current_gene and current_seq:
                        genes[current_gene] = "".join(current_seq).upper()
                    # Parse gene name from header
                    # Format: >gene_name_variant additional info
                    header = line[1:].split()[0]
                    current_gene = header
                    current_seq = []
                else:
                    current_seq.append(line)

            if current_gene and current_seq:
                genes[current_gene] = "".join(current_seq).upper()

        return genes

    def build_kmer_index(self) -> None:
        """Build k-mer index for all resistance genes."""
        logger.info(f"Building {self.kmer_size}-mer index for {len(self.resistance_genes)} genes...")

        for gene_name, gene_data in self.resistance_genes.items():
            seq = gene_data["sequence"]
            kmers = set()
            for i in range(len(seq) - self.kmer_size + 1):
                kmer = seq[i:i + self.kmer_size]
                if all(c in "ACGT" for c in kmer):
                    kmers.add(kmer)
                    # Also add reverse complement
                    kmers.add(self._reverse_complement(kmer))
            self.gene_kmers[gene_name] = kmers

        logger.info("K-mer index built successfully")

    @staticmethod
    def _reverse_complement(seq: str) -> str:
        """Get reverse complement of a DNA sequence."""
        complement = {"A": "T", "T": "A", "G": "C", "C": "G"}
        return "".join(complement.get(b, "N") for b in reversed(seq))

    def detect_genes_in_sequence(
        self,
        sequence: str,
        min_coverage: float = 0.5,
    ) -> List[Dict]:
        """Detect resistance genes in a genome sequence.

        Args:
            sequence: Genome sequence (DNA)
            min_coverage: Minimum fraction of gene k-mers that must be found

        Returns:
            List of detected genes with metadata
        """
        sequence = sequence.upper()

        # Build k-mer set for the genome
        genome_kmers = set()
        for i in range(len(sequence) - self.kmer_size + 1):
            kmer = sequence[i:i + self.kmer_size]
            if all(c in "ACGT" for c in kmer):
                genome_kmers.add(kmer)

        detected = []
        for gene_name, gene_kmers in self.gene_kmers.items():
            if not gene_kmers:
                continue

            # Count matching k-mers
            matching = len(gene_kmers & genome_kmers)
            coverage = matching / len(gene_kmers)

            if coverage >= min_coverage:
                gene_data = self.resistance_genes[gene_name]
                detected.append({
                    "gene": gene_name,
                    "drug_class": gene_data["drug_class"],
                    "coverage": round(coverage, 3),
                    "matching_kmers": matching,
                    "total_kmers": len(gene_kmers),
                })

        return detected

    def annotate_genomes(
        self,
        genomes_dir: str = "data/raw/ncbi/genomes",
        metadata_file: str = "data/raw/ncbi/complete_metadata.csv",
        output_dir: str = "data/raw/ncbi/amr_annotations",
        min_coverage: float = 0.5,
        max_genomes: Optional[int] = None,
    ) -> pd.DataFrame:
        """Annotate all genomes with AMR gene detection.

        Args:
            genomes_dir: Directory containing genome FASTA files
            metadata_file: Path to metadata CSV
            output_dir: Directory to save results
            min_coverage: Minimum k-mer coverage for gene detection
            max_genomes: Maximum number of genomes to process (for testing)

        Returns:
            DataFrame with AMR annotations for all genomes
        """
        genomes_path = Path(genomes_dir)
        output_path = Path(output_dir)
        output_path.mkdir(parents=True, exist_ok=True)

        # Load database and build index
        if not self.resistance_genes:
            self.load_resfinder_database()
        if not self.gene_kmers:
            self.build_kmer_index()

        # Get genome files
        genome_files = list(genomes_path.glob("*.fna.gz"))
        if max_genomes:
            genome_files = genome_files[:max_genomes]

        logger.info(f"Processing {len(genome_files)} genomes...")

        all_annotations = []
        genome_drug_classes = {}  # biosample_id -> set of drug classes

        for i, genome_file in enumerate(genome_files):
            biosample_id = genome_file.stem.replace(".fna", "")

            try:
                # Load genome sequence
                with gzip.open(genome_file, "rt") as f:
                    sequences = []
                    current_seq = []
                    for line in f:
                        line = line.strip()
                        if line.startswith(">"):
                            if current_seq:
                                sequences.append("".join(current_seq))
                                current_seq = []
                        else:
                            current_seq.append(line)
                    if current_seq:
                        sequences.append("".join(current_seq))

                    full_sequence = "".join(sequences)

                # Detect genes
                detected = self.detect_genes_in_sequence(full_sequence, min_coverage)

                # Record drug classes for this genome
                drug_classes_found = set()
                for gene_info in detected:
                    gene_info["biosample_id"] = biosample_id
                    all_annotations.append(gene_info)
                    drug_classes_found.add(gene_info["drug_class"])

                genome_drug_classes[biosample_id] = drug_classes_found

            except Exception as e:
                logger.warning(f"Error processing {biosample_id}: {e}")

            if (i + 1) % 50 == 0:
                logger.info(f"Processed {i + 1}/{len(genome_files)} genomes")

        # Create annotations DataFrame
        if all_annotations:
            annotations_df = pd.DataFrame(all_annotations)
            annotations_df.to_csv(output_path / "all_amr_annotations.csv", index=False)
            logger.info(f"Total AMR genes detected: {len(annotations_df)}")
        else:
            annotations_df = pd.DataFrame()
            logger.warning("No AMR genes detected in any genome")

        # Create labels DataFrame (multi-label format)
        labels_data = []
        for biosample_id, drug_classes in genome_drug_classes.items():
            row = {"biosample_id": biosample_id}
            for dc in sorted(self.drug_classes):
                row[dc] = 1 if dc in drug_classes else 0
            labels_data.append(row)

        if labels_data:
            labels_df = pd.DataFrame(labels_data)
            labels_df.to_csv(output_path / "amr_labels.csv", index=False)

            # Summary statistics
            logger.info("\nAMR Detection Summary:")
            logger.info(f"  Genomes processed: {len(genome_drug_classes)}")
            logger.info(f"  Genomes with AMR genes: {sum(1 for dc in genome_drug_classes.values() if dc)}")
            logger.info("\n  Resistance by drug class:")
            for dc in sorted(self.drug_classes):
                count = labels_df[dc].sum()
                logger.info(f"    {dc}: {count} genomes")
        else:
            labels_df = pd.DataFrame()

        # Save drug class mapping
        drug_mapping = {dc: i for i, dc in enumerate(sorted(self.drug_classes))}
        with open(output_path / "drug_class_mapping.json", "w") as f:
            json.dump(drug_mapping, f, indent=2)

        return labels_df

    def get_labels(self, output_dir: str = "data/raw/ncbi/amr_annotations") -> pd.DataFrame:
        """Load AMR labels from saved file."""
        labels_file = Path(output_dir) / "amr_labels.csv"
        if labels_file.exists():
            return pd.read_csv(labels_file)
        else:
            raise FileNotFoundError(
                f"Labels file not found: {labels_file}\n"
                "Run annotate_genomes() first."
            )


def main():
    """Run AMR gene detection on NCBI genomes."""
    print("=" * 60)
    print("AMR Gene Detection using ResFinder Database")
    print("=" * 60)

    detector = AMRGeneDetector(
        kmer_size=31,
        min_identity=0.8,
    )

    # Load database
    print("\nLoading ResFinder database...")
    detector.load_resfinder_database()

    # Build k-mer index
    print("\nBuilding k-mer index...")
    detector.build_kmer_index()

    # Annotate genomes
    print("\nAnnotating NCBI genomes...")
    labels_df = detector.annotate_genomes(
        min_coverage=0.5,
    )

    if len(labels_df) > 0:
        print("\n" + "=" * 60)
        print("Detection Complete!")
        print("=" * 60)
        print(f"\nResults saved to: data/raw/ncbi/amr_annotations/")
        print(f"  - amr_labels.csv: Multi-label resistance matrix")
        print(f"  - all_amr_annotations.csv: Detailed gene annotations")
        print(f"  - drug_class_mapping.json: Drug class to index mapping")
    else:
        print("\nNo AMR genes detected. Check genome files.")


if __name__ == "__main__":
    main()