src/preprocessing/patric_preprocessor.py

"""PATRIC database preprocessor for AMR prediction modeling."""

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

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder, StandardScaler

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


class PATRICPreprocessor:
    """Preprocess PATRIC data for AMR prediction models."""

    def __init__(
        self,
        patric_dir: str = "data/raw/patric",
        output_dir: str = "data/processed/patric",
    ):
        self.patric_dir = Path(patric_dir)
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        # Data containers
        self.amr_phenotypes: Optional[pd.DataFrame] = None
        self.genomes_metadata: Optional[pd.DataFrame] = None
        self.sequences: dict = {}
        self.label_encoders: dict = {}

    def load_data(self) -> None:
        """Load all PATRIC data files."""
        logger.info("Loading PATRIC data...")

        # Load AMR phenotypes
        amr_file = self.patric_dir / "amr_phenotypes.csv"
        if amr_file.exists():
            self.amr_phenotypes = pd.read_csv(amr_file)
            logger.info(f"Loaded {len(self.amr_phenotypes)} AMR phenotype records")
        else:
            raise FileNotFoundError(f"AMR phenotypes file not found: {amr_file}")

        # Load genome metadata
        meta_file = self.patric_dir / "genomes_metadata.csv"
        if meta_file.exists():
            self.genomes_metadata = pd.read_csv(meta_file)
            logger.info(f"Loaded {len(self.genomes_metadata)} genome metadata records")

        # Load genome sequences
        self._load_sequences()

    def _load_sequences(self) -> None:
        """Load genome sequences from FASTA files."""
        genomes_dir = self.patric_dir / "genomes"
        if not genomes_dir.exists():
            logger.warning(f"Genomes directory not found: {genomes_dir}")
            return

        fasta_files = list(genomes_dir.glob("*.fasta"))
        logger.info(f"Found {len(fasta_files)} genome FASTA files")

        for fasta_file in fasta_files:
            genome_id = fasta_file.stem
            sequences = []
            current_seq = []

            with open(fasta_file) as f:
                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))

            # Concatenate all contigs for this genome
            self.sequences[genome_id] = "".join(sequences)

        logger.info(f"Loaded sequences for {len(self.sequences)} genomes")

    def create_amr_dataset(
        self,
        antibiotic: Optional[str] = None,
        min_samples_per_class: int = 10,
    ) -> pd.DataFrame:
        """Create dataset mapping genomes to AMR phenotypes.

        Args:
            antibiotic: Specific antibiotic to filter. If None, uses all.
            min_samples_per_class: Minimum samples per class to include an antibiotic.

        Returns:
            DataFrame with genome_id, antibiotic, phenotype, and sequence.
        """
        if self.amr_phenotypes is None:
            self.load_data()

        df = self.amr_phenotypes.copy()

        # Filter to records with resistance phenotypes (Resistant/Susceptible)
        df = df[df["resistant_phenotype"].isin(["Resistant", "Susceptible"])].copy()
        logger.info(f"Records with R/S phenotypes: {len(df)}")

        # Filter by antibiotic if specified
        if antibiotic:
            df = df[df["antibiotic"] == antibiotic]
            logger.info(f"Records for {antibiotic}: {len(df)}")

        # Convert genome_id to string for matching with sequence keys
        df["genome_id"] = df["genome_id"].astype(str)

        # Add sequence data
        df["sequence"] = df["genome_id"].apply(
            lambda x: self.sequences.get(x, "")
        )

        # Filter to genomes with sequences
        df = df[df["sequence"].str.len() > 0].copy()
        logger.info(f"Records with genome sequences: {len(df)}")

        # Filter antibiotics with enough samples per class
        if not antibiotic:
            valid_antibiotics = []
            for ab in df["antibiotic"].unique():
                ab_df = df[df["antibiotic"] == ab]
                class_counts = ab_df["resistant_phenotype"].value_counts()
                if all(count >= min_samples_per_class for count in class_counts.values):
                    valid_antibiotics.append(ab)

            df = df[df["antibiotic"].isin(valid_antibiotics)]
            logger.info(f"Antibiotics with sufficient samples: {len(valid_antibiotics)}")
            logger.info(f"Valid antibiotics: {valid_antibiotics}")

        return df

    def extract_kmer_features(
        self, sequences: list, k: int = 6, max_features: int = 1000
    ) -> tuple:
        """Extract k-mer frequency features from DNA sequences.

        Args:
            sequences: List of DNA sequences.
            k: k-mer size (default 6 for DNA).
            max_features: Maximum number of k-mer features.

        Returns:
            Tuple of (feature_matrix, feature_names).
        """
        logger.info(f"Extracting {k}-mer features from {len(sequences)} sequences...")

        # Count all k-mers across sequences to find most common
        all_kmers = Counter()
        for seq in sequences:
            seq = seq.upper()
            for i in range(len(seq) - k + 1):
                kmer = seq[i : i + k]
                # Valid DNA nucleotides only
                if all(c in "ACGT" for c in kmer):
                    all_kmers[kmer] += 1

        # Select top k-mers as features
        top_kmers = [kmer for kmer, _ in all_kmers.most_common(max_features)]
        logger.info(f"Selected {len(top_kmers)} k-mer features")

        # Create feature matrix
        feature_matrix = np.zeros((len(sequences), len(top_kmers)))
        kmer_to_idx = {kmer: idx for idx, kmer in enumerate(top_kmers)}

        for seq_idx, seq in enumerate(sequences):
            seq = seq.upper()
            seq_len = len(seq) - k + 1
            if seq_len <= 0:
                continue
            for i in range(seq_len):
                kmer = seq[i : i + k]
                if kmer in kmer_to_idx:
                    feature_matrix[seq_idx, kmer_to_idx[kmer]] += 1
            # Normalize by sequence length
            if seq_len > 0:
                feature_matrix[seq_idx] /= seq_len

        return feature_matrix, top_kmers

    def prepare_single_antibiotic_data(
        self,
        antibiotic: str,
        k: int = 6,
        max_features: int = 1000,
        test_size: float = 0.2,
        val_size: float = 0.1,
        random_state: int = 42,
    ) -> dict:
        """Prepare dataset for a single antibiotic prediction.

        Args:
            antibiotic: Name of the antibiotic.
            k: k-mer size for feature extraction.
            max_features: Maximum number of k-mer features.
            test_size: Proportion for testing.
            val_size: Proportion of training data for validation.
            random_state: Random seed.

        Returns:
            Dictionary with train/val/test splits and metadata.
        """
        logger.info(f"Preparing data for {antibiotic} prediction...")

        # Create dataset for this antibiotic
        df = self.create_amr_dataset(antibiotic=antibiotic, min_samples_per_class=5)

        if len(df) < 20:
            raise ValueError(f"Not enough samples for {antibiotic}: {len(df)}")

        # Get unique genomes (one row per genome for this antibiotic)
        df_unique = df.drop_duplicates(subset=["genome_id"]).copy()
        logger.info(f"Unique genomes: {len(df_unique)}")

        # Extract features
        sequences = df_unique["sequence"].tolist()
        X, feature_names = self.extract_kmer_features(
            sequences, k=k, max_features=max_features
        )

        # Encode labels (binary: Resistant=1, Susceptible=0)
        le = LabelEncoder()
        y = le.fit_transform(df_unique["resistant_phenotype"])
        class_names = list(le.classes_)
        self.label_encoders[antibiotic] = le

        logger.info(f"Features shape: {X.shape}, Labels shape: {y.shape}")
        logger.info(f"Class distribution: {dict(zip(class_names, np.bincount(y)))}")

        # Split data
        try:
            X_temp, X_test, y_temp, y_test = train_test_split(
                X, y, test_size=test_size, random_state=random_state, stratify=y
            )
            val_ratio = val_size / (1 - test_size)
            X_train, X_val, y_train, y_val = train_test_split(
                X_temp, y_temp, test_size=val_ratio, random_state=random_state, stratify=y_temp
            )
        except ValueError as e:
            logger.warning(f"Stratified split failed ({e}), using random split")
            X_temp, X_test, y_temp, y_test = train_test_split(
                X, y, test_size=test_size, random_state=random_state
            )
            val_ratio = val_size / (1 - test_size)
            X_train, X_val, y_train, y_val = train_test_split(
                X_temp, y_temp, test_size=val_ratio, random_state=random_state
            )

        logger.info(f"Train: {X_train.shape[0]}, Val: {X_val.shape[0]}, Test: {X_test.shape[0]}")

        return {
            "X_train": X_train,
            "X_val": X_val,
            "X_test": X_test,
            "y_train": y_train,
            "y_val": y_val,
            "y_test": y_test,
            "feature_names": feature_names,
            "class_names": class_names,
            "task_type": "binary",
            "metadata": {
                "antibiotic": antibiotic,
                "k": k,
                "max_features": max_features,
                "n_samples": len(df_unique),
                "n_features": X.shape[1],
                "n_classes": len(class_names),
            },
        }

    def prepare_multi_antibiotic_data(
        self,
        antibiotics: Optional[list] = None,
        k: int = 6,
        max_features: int = 1000,
        test_size: float = 0.2,
        val_size: float = 0.1,
        random_state: int = 42,
    ) -> dict:
        """Prepare dataset for multi-label antibiotic resistance prediction.

        Args:
            antibiotics: List of antibiotics to include. If None, uses all valid ones.
            k: k-mer size for feature extraction.
            max_features: Maximum number of k-mer features.
            test_size: Proportion for testing.
            val_size: Proportion of training data for validation.
            random_state: Random seed.

        Returns:
            Dictionary with train/val/test splits and metadata.
        """
        logger.info("Preparing multi-antibiotic prediction data...")

        # Get valid dataset
        df = self.create_amr_dataset(min_samples_per_class=5)

        if antibiotics:
            df = df[df["antibiotic"].isin(antibiotics)]

        # Get unique antibiotics
        unique_antibiotics = sorted(df["antibiotic"].unique())
        logger.info(f"Using {len(unique_antibiotics)} antibiotics")

        # Create genome-level features and multi-label targets
        genome_ids = df["genome_id"].unique()
        logger.info(f"Unique genomes: {len(genome_ids)}")

        # Build genome-antibiotic resistance matrix
        genome_sequences = {}
        genome_labels = {gid: {} for gid in genome_ids}

        for _, row in df.iterrows():
            gid = row["genome_id"]
            ab = row["antibiotic"]
            phenotype = row["resistant_phenotype"]

            if gid not in genome_sequences:
                genome_sequences[gid] = row["sequence"]

            # Store resistance (1) or susceptible (0)
            genome_labels[gid][ab] = 1 if phenotype == "Resistant" else 0

        # Create feature matrix and label matrix
        sequences = [genome_sequences[gid] for gid in genome_ids]
        X, feature_names = self.extract_kmer_features(
            sequences, k=k, max_features=max_features
        )

        # Create multi-label target matrix
        y = np.zeros((len(genome_ids), len(unique_antibiotics)))
        for i, gid in enumerate(genome_ids):
            for j, ab in enumerate(unique_antibiotics):
                if ab in genome_labels[gid]:
                    y[i, j] = genome_labels[gid][ab]

        logger.info(f"Features shape: {X.shape}, Labels shape: {y.shape}")

        # Split data (can't stratify with multi-label)
        X_temp, X_test, y_temp, y_test = train_test_split(
            X, y, test_size=test_size, random_state=random_state
        )
        val_ratio = val_size / (1 - test_size)
        X_train, X_val, y_train, y_val = train_test_split(
            X_temp, y_temp, test_size=val_ratio, random_state=random_state
        )

        logger.info(f"Train: {X_train.shape[0]}, Val: {X_val.shape[0]}, Test: {X_test.shape[0]}")

        return {
            "X_train": X_train,
            "X_val": X_val,
            "X_test": X_test,
            "y_train": y_train,
            "y_val": y_val,
            "y_test": y_test,
            "feature_names": feature_names,
            "class_names": unique_antibiotics,
            "task_type": "multilabel",
            "metadata": {
                "antibiotics": unique_antibiotics,
                "k": k,
                "max_features": max_features,
                "n_samples": len(genome_ids),
                "n_features": X.shape[1],
                "n_classes": len(unique_antibiotics),
            },
        }

    def save_processed_data(self, data: dict, prefix: str = "patric") -> None:
        """Save processed data to disk."""
        logger.info(f"Saving processed data to {self.output_dir}")

        # Save numpy arrays
        np.save(self.output_dir / f"{prefix}_X_train.npy", data["X_train"])
        np.save(self.output_dir / f"{prefix}_X_val.npy", data["X_val"])
        np.save(self.output_dir / f"{prefix}_X_test.npy", data["X_test"])
        np.save(self.output_dir / f"{prefix}_y_train.npy", data["y_train"])
        np.save(self.output_dir / f"{prefix}_y_val.npy", data["y_val"])
        np.save(self.output_dir / f"{prefix}_y_test.npy", data["y_test"])

        # Save metadata
        metadata = {
            "feature_names": data["feature_names"],
            "class_names": data["class_names"],
            "task_type": data["task_type"],
            **data["metadata"],
        }
        with open(self.output_dir / f"{prefix}_metadata.json", "w") as f:
            json.dump(metadata, f, indent=2)

        logger.info("Data saved successfully!")

    def get_antibiotic_statistics(self) -> pd.DataFrame:
        """Get statistics about antibiotics in the dataset."""
        if self.amr_phenotypes is None:
            self.load_data()

        # Filter to R/S phenotypes
        df = self.amr_phenotypes[
            self.amr_phenotypes["resistant_phenotype"].isin(["Resistant", "Susceptible"])
        ].copy()

        # Count by antibiotic
        stats = []
        for ab in df["antibiotic"].unique():
            ab_df = df[df["antibiotic"] == ab]
            r_count = (ab_df["resistant_phenotype"] == "Resistant").sum()
            s_count = (ab_df["resistant_phenotype"] == "Susceptible").sum()
            total = r_count + s_count
            r_ratio = r_count / total if total > 0 else 0

            stats.append({
                "antibiotic": ab,
                "resistant": r_count,
                "susceptible": s_count,
                "total": total,
                "resistance_rate": round(r_ratio, 3),
            })

        stats_df = pd.DataFrame(stats)
        stats_df = stats_df.sort_values("total", ascending=False)
        return stats_df

    def compute_gc_content(self, sequence: str) -> float:
        """Calculate GC content of a DNA sequence.

        Args:
            sequence: DNA sequence string.

        Returns:
            GC content as a fraction (0-1).
        """
        sequence = sequence.upper()
        if len(sequence) == 0:
            return 0.0
        gc_count = sequence.count("G") + sequence.count("C")
        return gc_count / len(sequence)

    def get_sequence_statistics(self) -> pd.DataFrame:
        """Get statistics about genome sequences.

        Returns:
            DataFrame with sequence statistics per genome.
        """
        if not self.sequences:
            self._load_sequences()

        stats = []
        for genome_id, seq in self.sequences.items():
            seq_len = len(seq)
            gc_content = self.compute_gc_content(seq)

            # Count nucleotides
            seq_upper = seq.upper()
            a_count = seq_upper.count("A")
            t_count = seq_upper.count("T")
            g_count = seq_upper.count("G")
            c_count = seq_upper.count("C")
            n_count = seq_upper.count("N")
            other_count = seq_len - a_count - t_count - g_count - c_count - n_count

            stats.append({
                "genome_id": genome_id,
                "length": seq_len,
                "gc_content": round(gc_content, 4),
                "a_count": a_count,
                "t_count": t_count,
                "g_count": g_count,
                "c_count": c_count,
                "n_count": n_count,
                "other_count": other_count,
            })

        return pd.DataFrame(stats).sort_values("length", ascending=False)

    def get_organism_statistics(self) -> pd.DataFrame:
        """Get statistics about organisms in the dataset.

        Returns:
            DataFrame with organism-level statistics.
        """
        if self.genomes_metadata is None:
            self.load_data()

        if self.genomes_metadata is None:
            logger.warning("No genome metadata available")
            return pd.DataFrame()

        # Group by species/organism
        if "species" in self.genomes_metadata.columns:
            group_col = "species"
        elif "organism_name" in self.genomes_metadata.columns:
            group_col = "organism_name"
        else:
            logger.warning("No species or organism_name column found")
            return pd.DataFrame()

        stats = self.genomes_metadata.groupby(group_col).agg({
            "genome_id": "count",
        }).reset_index()
        stats.columns = [group_col, "genome_count"]
        stats = stats.sort_values("genome_count", ascending=False)

        return stats

    def get_phenotype_by_organism(self) -> pd.DataFrame:
        """Get resistance statistics broken down by organism.

        Returns:
            DataFrame with resistance rates per organism and antibiotic.
        """
        if self.amr_phenotypes is None or self.genomes_metadata is None:
            self.load_data()

        # Merge phenotypes with metadata
        df = self.amr_phenotypes.merge(
            self.genomes_metadata[["genome_id", "species"]].drop_duplicates(),
            on="genome_id",
            how="left"
        )

        # Filter to R/S phenotypes
        df = df[df["resistant_phenotype"].isin(["Resistant", "Susceptible"])].copy()

        # Calculate resistance rate per organism-antibiotic pair
        stats = []
        for (species, ab), group in df.groupby(["species", "antibiotic"]):
            r_count = (group["resistant_phenotype"] == "Resistant").sum()
            s_count = (group["resistant_phenotype"] == "Susceptible").sum()
            total = r_count + s_count
            r_rate = r_count / total if total > 0 else 0

            stats.append({
                "species": species,
                "antibiotic": ab,
                "resistant": r_count,
                "susceptible": s_count,
                "total": total,
                "resistance_rate": round(r_rate, 3),
            })

        return pd.DataFrame(stats).sort_values(["species", "total"], ascending=[True, False])

    def extract_combined_features(
        self,
        sequences: List[str],
        k: int = 6,
        max_features: int = 1000,
        include_gc: bool = True,
        include_length: bool = True,
    ) -> Tuple[np.ndarray, List[str]]:
        """Extract k-mer features combined with sequence statistics.

        Args:
            sequences: List of DNA sequences.
            k: k-mer size.
            max_features: Maximum number of k-mer features.
            include_gc: Whether to include GC content feature.
            include_length: Whether to include normalized sequence length.

        Returns:
            Tuple of (feature_matrix, feature_names).
        """
        # Get k-mer features
        kmer_features, kmer_names = self.extract_kmer_features(
            sequences, k=k, max_features=max_features
        )

        feature_names = kmer_names.copy()
        additional_features = []

        if include_gc:
            gc_features = np.array([self.compute_gc_content(seq) for seq in sequences])
            additional_features.append(gc_features.reshape(-1, 1))
            feature_names.append("gc_content")

        if include_length:
            lengths = np.array([len(seq) for seq in sequences])
            # Log-normalize length
            log_lengths = np.log1p(lengths)
            # Scale to 0-1 range
            if log_lengths.max() > log_lengths.min():
                log_lengths = (log_lengths - log_lengths.min()) / (log_lengths.max() - log_lengths.min())
            additional_features.append(log_lengths.reshape(-1, 1))
            feature_names.append("log_length_normalized")

        if additional_features:
            additional_matrix = np.hstack(additional_features)
            feature_matrix = np.hstack([kmer_features, additional_matrix])
        else:
            feature_matrix = kmer_features

        logger.info(f"Combined features shape: {feature_matrix.shape}")
        return feature_matrix, feature_names

    def get_data_summary(self) -> Dict:
        """Get comprehensive summary of the PATRIC dataset.

        Returns:
            Dictionary containing dataset summary statistics.
        """
        if self.amr_phenotypes is None:
            self.load_data()

        summary = {
            "total_amr_records": len(self.amr_phenotypes),
            "total_genomes_with_sequences": len(self.sequences),
            "unique_antibiotics": self.amr_phenotypes["antibiotic"].nunique(),
            "antibiotics_list": sorted(self.amr_phenotypes["antibiotic"].unique().tolist()),
        }

        # R/S phenotype breakdown
        rs_df = self.amr_phenotypes[
            self.amr_phenotypes["resistant_phenotype"].isin(["Resistant", "Susceptible"])
        ]
        summary["resistant_records"] = int((rs_df["resistant_phenotype"] == "Resistant").sum())
        summary["susceptible_records"] = int((rs_df["resistant_phenotype"] == "Susceptible").sum())
        summary["records_with_rs_phenotype"] = len(rs_df)

        # Genome metadata
        if self.genomes_metadata is not None:
            summary["total_genome_metadata_records"] = len(self.genomes_metadata)
            if "species" in self.genomes_metadata.columns:
                summary["unique_species"] = self.genomes_metadata["species"].nunique()
                summary["species_list"] = sorted(
                    self.genomes_metadata["species"].dropna().unique().tolist()
                )

        # Sequence statistics
        if self.sequences:
            seq_lengths = [len(seq) for seq in self.sequences.values()]
            summary["sequence_stats"] = {
                "count": len(seq_lengths),
                "min_length": min(seq_lengths),
                "max_length": max(seq_lengths),
                "mean_length": int(np.mean(seq_lengths)),
                "median_length": int(np.median(seq_lengths)),
            }

        return summary


def main():
    """Main preprocessing pipeline."""
    preprocessor = PATRICPreprocessor()

    # Load data
    preprocessor.load_data()

    # Show statistics
    print("\n=== Antibiotic Statistics ===")
    stats = preprocessor.get_antibiotic_statistics()
    print(stats.head(20).to_string(index=False))

    # Prepare multi-antibiotic prediction data
    print("\n=== Preparing Multi-Antibiotic Prediction Data ===")
    try:
        data = preprocessor.prepare_multi_antibiotic_data(
            k=6,
            max_features=500,
            test_size=0.2,
            val_size=0.1,
        )
        preprocessor.save_processed_data(data, prefix="patric_multilabel")
        print("Saved multi-label prediction data")
    except Exception as e:
        print(f"Error preparing multi-antibiotic data: {e}")

    # Prepare single antibiotic models for top antibiotics
    print("\n=== Preparing Single-Antibiotic Models ===")
    top_antibiotics = stats.head(5)["antibiotic"].tolist()

    for ab in top_antibiotics:
        try:
            print(f"\nProcessing: {ab}")
            data = preprocessor.prepare_single_antibiotic_data(
                antibiotic=ab,
                k=6,
                max_features=500,
            )
            # Clean antibiotic name for filename
            ab_clean = ab.replace("/", "_").replace(" ", "_").lower()
            preprocessor.save_processed_data(data, prefix=f"patric_{ab_clean}")
            print(f"Saved {ab} prediction data")
        except Exception as e:
            print(f"Error preparing {ab} data: {e}")

    print("\n=== Preprocessing Complete ===")
    print(f"Output directory: {preprocessor.output_dir}")


if __name__ == "__main__":
    main()