src/ml/feature_extraction.py

"""Feature extraction for real-time AMR prediction.

This module extracts k-mer features from DNA sequences for prediction.
Uses the same k-mer vocabulary as the trained model.
"""

import json
import gzip
import logging
from pathlib import Path
from typing import List, Tuple, Optional
from collections import Counter

import numpy as np

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

# Default paths
PROJECT_ROOT = Path(__file__).parent.parent.parent
METADATA_PATH = PROJECT_ROOT / "data" / "processed" / "ncbi" / "ncbi_amr_metadata.json"


class KmerFeatureExtractor:
    """Extract k-mer features from DNA sequences using trained vocabulary.

    This extractor uses the same k-mer vocabulary that was used during model
    training to ensure consistent feature extraction for inference.
    """

    def __init__(self, metadata_path: Optional[str] = None):
        """Initialize the feature extractor.

        Args:
            metadata_path: Path to metadata JSON containing feature_names.
                          If None, uses default path.
        """
        self.metadata_path = Path(metadata_path) if metadata_path else METADATA_PATH
        self.feature_names: List[str] = []
        self.k: int = 6
        self.kmer_to_idx: dict = {}

        self._load_vocabulary()

    def _load_vocabulary(self):
        """Load k-mer vocabulary from metadata file."""
        if not self.metadata_path.exists():
            raise FileNotFoundError(f"Metadata file not found: {self.metadata_path}")

        with open(self.metadata_path) as f:
            metadata = json.load(f)

        self.feature_names = metadata.get("feature_names", [])
        self.k = metadata.get("k", 6)

        if not self.feature_names:
            raise ValueError("No feature_names found in metadata")

        self.kmer_to_idx = {kmer: idx for idx, kmer in enumerate(self.feature_names)}

        logger.info(f"Loaded {len(self.feature_names)} k-mer features (k={self.k})")

    def extract_features(self, sequence: str) -> np.ndarray:
        """Extract k-mer features from a single DNA sequence.

        Args:
            sequence: DNA sequence string (A, C, G, T characters)

        Returns:
            Feature vector of shape (n_features,) with k-mer frequencies
        """
        sequence = sequence.upper().replace('\n', '').replace(' ', '')
        seq_len = len(sequence) - self.k + 1

        if seq_len <= 0:
            logger.warning(f"Sequence too short (len={len(sequence)}, need >={self.k})")
            return np.zeros(len(self.feature_names))

        # Count k-mers
        features = np.zeros(len(self.feature_names))
        valid_count = 0

        for i in range(seq_len):
            kmer = sequence[i:i + self.k]
            # Only count valid DNA k-mers
            if all(c in "ACGT" for c in kmer):
                valid_count += 1
                if kmer in self.kmer_to_idx:
                    features[self.kmer_to_idx[kmer]] += 1

        # Normalize by total valid k-mers
        if valid_count > 0:
            features = features / valid_count

        return features

    def extract_features_batch(self, sequences: List[str]) -> np.ndarray:
        """Extract k-mer features from multiple sequences.

        Args:
            sequences: List of DNA sequence strings

        Returns:
            Feature matrix of shape (n_sequences, n_features)
        """
        return np.array([self.extract_features(seq) for seq in sequences])

    def parse_fasta(self, content: str) -> List[Tuple[str, str]]:
        """Parse FASTA format content.

        Args:
            content: FASTA file content as string

        Returns:
            List of (header, sequence) tuples
        """
        sequences = []
        current_header = None
        current_seq = []

        for line in content.strip().split('\n'):
            line = line.strip()
            if line.startswith('>'):
                if current_header is not None:
                    sequences.append((current_header, ''.join(current_seq)))
                current_header = line[1:]
                current_seq = []
            else:
                current_seq.append(line)

        if current_header is not None:
            sequences.append((current_header, ''.join(current_seq)))

        return sequences

    def parse_fastq(self, content: str) -> List[Tuple[str, str]]:
        """Parse FASTQ format content.

        Args:
            content: FASTQ file content as string

        Returns:
            List of (header, sequence) tuples
        """
        sequences = []
        lines = content.strip().split('\n')

        i = 0
        while i < len(lines):
            if lines[i].startswith('@'):
                header = lines[i][1:]
                sequence = lines[i + 1] if i + 1 < len(lines) else ''
                sequences.append((header, sequence))
                i += 4  # Skip quality lines
            else:
                i += 1

        return sequences

    def extract_from_file_content(
        self,
        content: str,
        file_format: str = "fasta"
    ) -> Tuple[np.ndarray, List[str]]:
        """Extract features from file content.

        Args:
            content: File content as string
            file_format: Either 'fasta' or 'fastq'

        Returns:
            Tuple of (feature_matrix, sequence_headers)
        """
        if file_format.lower() in ['fastq', 'fq']:
            sequences = self.parse_fastq(content)
        else:
            sequences = self.parse_fasta(content)

        if not sequences:
            raise ValueError("No sequences found in file content")

        headers = [h for h, _ in sequences]
        seqs = [s for _, s in sequences]

        # For multiple sequences, concatenate them (typical for assembled genomes)
        if len(seqs) > 1:
            logger.info(f"Found {len(seqs)} sequences, concatenating for feature extraction")
            combined_seq = ''.join(seqs)
            features = self.extract_features(combined_seq)
            return features.reshape(1, -1), [f"combined_{len(seqs)}_sequences"]
        else:
            features = self.extract_features(seqs[0])
            return features.reshape(1, -1), headers

    @property
    def n_features(self) -> int:
        """Number of features (k-mers) in vocabulary."""
        return len(self.feature_names)


# Global extractor instance (lazy loaded)
_extractor: Optional[KmerFeatureExtractor] = None


def get_extractor() -> KmerFeatureExtractor:
    """Get or create the global feature extractor instance."""
    global _extractor
    if _extractor is None:
        _extractor = KmerFeatureExtractor()
    return _extractor


def extract_features_from_sequence(sequence: str) -> np.ndarray:
    """Convenience function to extract features from a sequence.

    Args:
        sequence: DNA sequence string

    Returns:
        Feature vector of shape (n_features,)
    """
    return get_extractor().extract_features(sequence)


def extract_features_from_fasta(content: str) -> np.ndarray:
    """Convenience function to extract features from FASTA content.

    Args:
        content: FASTA file content

    Returns:
        Feature vector of shape (n_features,)
    """
    features, _ = get_extractor().extract_from_file_content(content, "fasta")
    return features[0]  # Return first (and typically only) row