import os
import time
import csv
import yaml
import shutil
import subprocess
from snakemake.io import expand
import pandas as pd
import glob
import sys
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor
from threading import Lock
from pathlib import Path


# Name of the MitoGeneExtractor executable provided on PATH by the BeeGees conda
# environment (the bioconda 'mitogeneextractor' package). MGE is no longer
# referenced via a user-supplied 'mge_path' in the config; it is resolved from
# the active conda environment instead. If the bioconda package exposes the
# command under a different name, update this single constant.
MGE_EXECUTABLE = "MitoGeneExtractor"


# ===== CONFIGURATION VALIDATION =====
def validate_config(config):
    """Validate all configuration parameters"""
    required_keys = ['samples_file', 'output_dir', 'barcode_recovery', 'run_name']
    missing_keys = [key for key in required_keys if key not in config]

    if missing_keys:
        sys.stderr.write(f"WARNING: Missing required configuration keys: {', '.join(missing_keys)}\n")
        sys.exit(1)

    # Validate file existence
    if not os.path.exists(config['samples_file']):
        sys.stderr.write(f"WARNING: Samples file not found: {config['samples_file']}\n")
        sys.exit(1)

    # Validate MitoGeneExtractor is available on PATH (installed via the
    # BeeGees conda environment, i.e. the bioconda 'mitogeneextractor' package).
    if shutil.which(MGE_EXECUTABLE) is None:
        sys.stderr.write(
            f"WARNING: MitoGeneExtractor executable '{MGE_EXECUTABLE}' not found on PATH. "
            f"Ensure the BeeGees conda environment (with the 'mitogeneextractor' package) is activated.\n")
        sys.exit(1)

    # Validate barcode_recovery parameters
    br = config.get('barcode_recovery', {})
    if not isinstance(br.get('r'), list) or not br.get('r'):
        sys.stderr.write("WARNING: 'barcode_recovery.r' must be a non-empty list\n")
        sys.exit(1)

    if not isinstance(br.get('s'), list) or not br.get('s'):
        sys.stderr.write("WARNING: 'barcode_recovery.s' must be a non-empty list\n")
        sys.exit(1)

def validate_fastp_config(config):
    """Validate fastp configuration parameters"""
    fastp_config = config.get("fastp", {})

    # Validate adapter sequences if provided - must be non-empty strings
    adapter_r1 = fastp_config.get("adapter_r1", "")
    adapter_r2 = fastp_config.get("adapter_r2", "")

    if adapter_r1 is not None and not isinstance(adapter_r1, str):
        sys.stderr.write("WARNING: 'fastp.adapter_r1' must be a string\n")
        sys.exit(1)

    if adapter_r2 is not None and not isinstance(adapter_r2, str):
        sys.stderr.write("WARNING: 'fastp.adapter_r2' must be a string\n")
        sys.exit(1)

    # If one adapter is provided, both should be provided
    if bool(adapter_r1) != bool(adapter_r2):
        sys.stderr.write("WARNING: Both 'fastp.adapter_r1' and 'fastp.adapter_r2' must be provided together, or neither\n")
        sys.exit(1)

    # Validate extra_fastp_args is a string if provided
    extra_args = fastp_config.get("extra_fastp_args", "")
    if extra_args is not None and not isinstance(extra_args, str):
        sys.stderr.write("WARNING: 'fastp.extra_fastp_args' must be a string\n")
        sys.exit(1)

def validate_gene_fetch_config(config):
    """Validate gene_fetch specific configuration"""
    run_gene_fetch = config.get("run_gene_fetch", True)
    if not isinstance(run_gene_fetch, bool):
        sys.stderr.write("WARNING: 'run_gene_fetch' must be a boolean (true/false)\n")
        sys.exit(1)
    
    if run_gene_fetch:
        gene_fetch_config = config.get("gene_fetch", {})
        required_gene_fetch_keys = ['email', 'api_key', 'gene', 'input_type', 'minimum_length']
        missing_gene_fetch_keys = [key for key in required_gene_fetch_keys 
                                 if key not in gene_fetch_config or not gene_fetch_config[key]]
        
        if missing_gene_fetch_keys:
            sys.stderr.write(f"WARNING: Missing required gene_fetch configuration parameters: {', '.join(missing_gene_fetch_keys)}\n")
            sys.exit(1)
            
        # Validate input_type value
        input_type = gene_fetch_config.get("input_type", "").lower()
        if input_type not in ["taxid", "hierarchical"]:
            sys.stderr.write("WARNING: 'input_type' must be either 'taxid' or 'hierarchical'\n")
            sys.exit(1)
        
        # Validate minimum_length value
        minimum_length = gene_fetch_config.get("minimum_length", 500)
        try:
            minimum_length = int(minimum_length)
        except (ValueError, TypeError):
            sys.stderr.write("ERROR: 'minimum_length' must be an integer\n")
            sys.exit(1)
            
        # Validate samples_file exists
        samples_file = gene_fetch_config.get("samples_file", config.get("samples_file", ""))
        if not samples_file or not os.path.exists(samples_file):
            sys.stderr.write(f"WARNING: Gene fetch samples file not found: {samples_file}\n")
            sys.exit(1)
    else:
        # If not using gene_fetch, validate sequence_reference_file exists
        if 'sequence_reference_file' not in config:
            sys.stderr.write("WARNING: 'sequence_reference_file' is required when run_gene_fetch is false\n")
            sys.exit(1)
        
        if not os.path.exists(config['sequence_reference_file']):
            sys.stderr.write(f"WARNING: Sequence reference file not found: {config['sequence_reference_file']}\n")
            sys.exit(1)

def validate_structural_validation_config(config):
    """Validate structural validation configuration"""
    sv = config.get("structural_validation", {})
    if not sv.get("target"):
        sys.stderr.write("WARNING: 'structural_validation.target' is required (e.g. 'cox1' or 'rbcl')\n")
        sys.exit(1)

def validate_taxonomic_validation_config(config):
    """Validate taxonomic validation configuration"""
    taxonomic_validation_params = config.get("taxonomic_validation", {})

    # Validate blast_db path
    database_path = taxonomic_validation_params.get("blast_db", "")
    if not database_path:
        sys.stderr.write("WARNING: 'taxonomic_validation.blast_db' path is required\n")
        sys.exit(1)

    if not os.path.exists(database_path):
        sys.stderr.write(f"WARNING: Taxonomic validation (BLASTn) database not found: {database_path}\n")
        sys.exit(1)

    # Validate taxonomy file
    db_taxonomy = taxonomic_validation_params.get("db_taxonomy", "")
    if not db_taxonomy:
        sys.stderr.write("WARNING: 'taxonomic_validation.db_taxonomy' path is required\n")
        sys.exit(1)
    if not os.path.exists(db_taxonomy):
        sys.stderr.write(f"WARNING: Taxonomy file not found: {db_taxonomy}\n")
        sys.exit(1)

    # Validate expected taxonomy file
    expected_taxonomy = taxonomic_validation_params.get("expected_taxonomy", "")
    if not expected_taxonomy:
        sys.stderr.write("WARNING: 'taxonomic_validation.expected_taxonomy' path is required\n")
        sys.exit(1)
    if not os.path.exists(expected_taxonomy):
        sys.stderr.write(f"WARNING: Expected taxonomy file not found: {expected_taxonomy}\n")
        sys.exit(1)
        
        # Validate min_pident parameter
        min_pident = taxonomic_validation_params.get("min_pident", "80")
        try:
            min_pident_float = float(min_pident)
            if not (0 <= min_pident_float <= 100):
                sys.stderr.write("WARNING: 'min_pident' must be between 0 and 100\n")
                sys.exit(1)
        except (ValueError, TypeError):
            sys.stderr.write("WARNING: 'min_pident' must be a valid number\n")
            sys.exit(1)

# Run all validations
validate_config(config)
validate_fastp_config(config)
validate_gene_fetch_config(config)
validate_structural_validation_config(config)
validate_taxonomic_validation_config(config)




# ===== DATA PARSING FUNCTIONS =====
def parse_samples(samples_file):
    """Parse samples from .csv files"""
    samples = {}
    with open(samples_file, mode='r') as infile:
        reader = csv.DictReader(infile)
        
        # Get the actual column names from the CSV
        fieldnames = reader.fieldnames
        
        # Check for required columns with helpful error messages.
        # 'reverse' is OPTIONAL: a single-end (SE/Ultima) sample sheet may omit
        # the R2 column entirely, or include it with empty values.
        required_columns = {
            'ID': ['ID', 'id', 'sample_id', 'SampleID'],
            'forward': ['forward', 'fwd', 'R1', 'read1'],
        }
        optional_columns = {
            'reverse': ['reverse', 'rev', 'R2', 'read2'],
        }
        
        column_mapping = {}
        missing_columns = []
        
        for required, alternatives in required_columns.items():
            found = False
            for alt in alternatives:
                if alt in fieldnames:
                    column_mapping[required] = alt
                    found = True
                    break
            if not found:
                missing_columns.append(f"{required} (acceptable alternatives: {', '.join(alternatives)})")
        
        if missing_columns:
            sys.stderr.write(f"ERROR: samples.csv is missing required columns:\n")
            for missing in missing_columns:
                sys.stderr.write(f"  - {missing}\n")
            sys.stderr.write(f"\nFound columns in CSV: {', '.join(fieldnames)}\n")
            sys.stderr.write(f"\nPlease ensure your CSV has columns for sample ID and forward/reverse reads.\n")
            sys.exit(1)

        # Optional reverse column (absent => single-end run)
        for optional, alternatives in optional_columns.items():
            for alt in alternatives:
                if alt in fieldnames:
                    column_mapping[optional] = alt
                    break
        
        # Parse rows using the mapped column names
        for row in reader:
            sample_id = row[column_mapping['ID']]
            forward_read = row[column_mapping['forward']]
            if 'reverse' in column_mapping:
                reverse_read = row[column_mapping['reverse']]
            else:
                reverse_read = ""
            samples[sample_id] = {"R1": forward_read, "R2": reverse_read}
    
    return samples

def parse_sequence_references(sequence_reference_file):
    """Parse references from CSV"""
    sequence_references = {}
    with open(sequence_reference_file, mode='r') as infile:
        reader = csv.DictReader(infile)
        for row in reader:
            sample_id = row['ID']
            protein_reference_path = row['protein_reference_path']
            sequence_references[sample_id] = {
                "protein_path": protein_reference_path,
            }    
    return sequence_references



# ===== CONFIGURATION PARAMETERS =====
# Path configuration
samples = parse_samples(config["samples_file"])
run_name = config["run_name"]

# Fastp configuration
# Default adapters (Illumina TruSeq) are used if not specified in config.
_fastp_config = config.get("fastp", {})
_default_adapter_r1 = "AGATCGGAAGAGCACACGTCTGAACTCCAGTCA"
_default_adapter_r2 = "AGATCGGAAGAGCGTCGTGTAGGGAAAGAGTGT"
fastp_adapter_r1 = _fastp_config.get("adapter_r1", "") or _default_adapter_r1
fastp_adapter_r2 = _fastp_config.get("adapter_r2", "") or _default_adapter_r2
# Single extra-args field applied to both PE and SE fastp runs.
fastp_extra_args = _fastp_config.get("extra_fastp_args", "") or ""

# Gene fetch configuration
run_gene_fetch = config.get("run_gene_fetch", True)
if run_gene_fetch:
    sequence_reference_file_path = os.path.join(config["output_dir"], "02_references", "sequence_references.csv")
else:
    sequence_reference_file_path = config["sequence_reference_file"]

# Downsampling configuration
downsampling_enabled = config.get("downsampling", {}).get("enabled", False)
max_reads = config.get("downsampling", {}).get("max_reads", 0)

# MGE / barcode recovery parameter configuration
_br = config["barcode_recovery"]
r = _br["r"]
s = _br["s"]
n = _br["n"]
C = _br["C"]
t = _br["t"]

# Fasta cleaner configuration
fasta_cleaner_params = config.get("fasta_cleaner", {
    "consensus_threshold": 0.5,
    "human_threshold": 0.95,
    "at_difference": 0.1,
    "at_mode": "absolute",
    "outlier_percentile": 90.0
})

# Structural validation configuration
structural_validation_params = config.get("structural_validation", {
    "target": "cox1",
    "verbose": False,
    "relaxed": False
})

# Taxonomic validation configuration
# NOTE: the fallback dict below is unreachable in practice -
# validate_taxonomic_validation_config() runs earlier and exits when the
# taxonomic_validation key is absent or blast_db is empty. It is kept only so
# this module is importable in isolation. BLAST search options are not
# configurable here: they are the defaults in scripts/tv_local_blast.py, and the
# rule does not pass --blast-opts.
taxonomic_validation_params = config.get("taxonomic_validation", {
    "blast_db": "",
    "verbose": True,
    "db_taxonomy": "",
    "expected_taxonomy": "",
    "taxval_rank": "family",
    "min_pident": "80"
})

# Rule resources
rule_resources = config["rules"]


# ===== RUN MODE DETECTION (uniform SE vs PE) =====
# A run is uniform: either all samples carry an R2 value (PE) or none do
# (single-end / Ultima). A mixed sheet is a user input error.
has_r2 = [bool(str(s.get("R2", "")).strip()) for s in samples.values()]
if any(has_r2) and not all(has_r2):
    sys.stderr.write(
        "ERROR: Mixed SE/PE samples detected in samples.csv. "
        "All samples must have either R1 only (SE/Ultima) "
        "or both R1 and R2 (PE).\n"
    )
    sys.exit(1)

run_mode = "PE" if all(has_r2) else "SE"

if run_mode == "SE":
    overrep_fasta_path = os.path.join(workflow.basedir, "..", "resources/overrep_fasta/overrep.fasta")
    if not os.path.exists(overrep_fasta_path):
        sys.stderr.write(f"WARNING: SE adapter fasta not found: {overrep_fasta_path}\n")
        sys.exit(1)


# ===== TAXDUMP DIRECTORY =====
# Defaults to resources/ncbi_taxdump/ in the user's working directory so that
# when BeeGees is installed as a PyPI package the workflow can still write to a
# writable location (site-packages is read-only). Override via taxdump_dir in
# config.yaml when a shared/pre-downloaded dump is preferred.
taxdump_dir = config.get("taxdump_dir", os.path.join(os.getcwd(), "resources", "ncbi_taxdump"))


# ===== DIRECTORY STRUCTURE SETUP =====
# Create main output directory
main_output_dir = config["output_dir"]
try:
    Path(main_output_dir).mkdir(parents=True, exist_ok=True)
except PermissionError:
    sys.exit(
        f"ERROR: Cannot create output directory '{main_output_dir}'.\n"
        f"  Check that 'output_dir' in your config is set to a writable path."
    )

# Create preprocessing directories
preprocessing_dir = os.path.join(main_output_dir, "01_preprocessing")
preprocessing_dir_merge = os.path.join(preprocessing_dir, "merge_mode")
preprocessing_dir_concat = os.path.join(preprocessing_dir, "concat_mode")
preprocessing_dir_se = os.path.join(preprocessing_dir, "se_mode")

# Create barcode recovery directories
barcode_recovery_dir = os.path.join(main_output_dir, "03_barcode_recovery")
barcode_recovery_dir_merge = os.path.join(barcode_recovery_dir, "merge_mode")
barcode_recovery_dir_concat = os.path.join(barcode_recovery_dir, "concat_mode")
barcode_recovery_dir_se = os.path.join(barcode_recovery_dir, "se_mode")

# Create references directory if using gene_fetch
if run_gene_fetch:
    references_dir = os.path.join(main_output_dir, "02_references")
    Path(references_dir).mkdir(parents=True, exist_ok=True)



# ===== HELPER FUNCTIONS =====
def get_hmm_file_for_target(target):
    """Map target gene to corresponding HMM file path"""
    hmm_mapping = {
        "cox1": os.path.join(workflow.basedir, "..", "resources/hmm/COI-5P.hmm"),
        "coi": os.path.join(workflow.basedir, "..", "resources/hmm/COI-5P.hmm"),
        "rbcl": os.path.join(workflow.basedir, "..", "resources/hmm/rbcL.hmm")
    }
    
    target_lower = target.lower()
    if target_lower not in hmm_mapping:
        raise ValueError(f"Unknown target '{target}'. Available targets: {list(hmm_mapping.keys())}")
    
    return hmm_mapping[target_lower]

def get_protein_reference(wildcards):
    """Get protein reference path for a sample"""
    if run_gene_fetch:
        return os.path.join(main_output_dir, f"02_references/protein/{wildcards.sample}.fasta")
    else:
        sequence_refs = parse_sequence_references(config["sequence_reference_file"])
        return sequence_refs[wildcards.sample]["protein_path"]

def get_all_mge_consensus_files(mode):
    """Generate list of expected consensus files for a given mode"""
    if mode == "merge":
        output_dir = barcode_recovery_dir_merge
    elif mode == "se":
        output_dir = barcode_recovery_dir_se
    else:
        output_dir = barcode_recovery_dir_concat
    files = []
    
    for sample in samples.keys():
        for r_val in r:
            for s_val in s:
                files.append(os.path.join(output_dir, f"consensus/{sample}_r_{r_val}_s_{s_val}_con_{sample}.fas"))
    
    return files

def get_mge_input_merge(wildcards):
    """Get the right input for merge mode"""
    return os.path.join(preprocessing_dir_merge, f"trimmed_data/{wildcards.sample}/{wildcards.sample}_merged_clean.fastq")

def get_mge_input_concat(wildcards):
    """Get input for concat mode"""
    if downsampling_enabled:
        return os.path.join(preprocessing_dir_concat, f"trimmed_data/{wildcards.sample}/{wildcards.sample}_concat_trimmed_downsampled.fastq")
    else:
        return os.path.join(preprocessing_dir_concat, f"trimmed_data/{wildcards.sample}/{wildcards.sample}_concat_trimmed.fq")

def get_mge_input_se(wildcards):
    """Get input for single-end (SE/Ultima) mode"""
    if downsampling_enabled:
        return os.path.join(preprocessing_dir_se, f"trimmed_data/{wildcards.sample}/{wildcards.sample}_se_trimmed_downsampled.fastq")
    else:
        return os.path.join(preprocessing_dir_se, f"trimmed_data/{wildcards.sample}/{wildcards.sample}_se_trimmed.fastq")

# Set up derived parameters
hmm_file_path = get_hmm_file_for_target(structural_validation_params["target"])
reference_dir = fasta_cleaner_params.get("reference_dir", None)
use_reference_filtering = reference_dir and reference_dir not in ["None", "null", "", None]



# ===== SNAKEMAKE RULE ALL =====
# Targets are grouped so PE (merge+concat) and SE (Ultima) runs request only the
# outputs their active path produces. Validation, reporting and final-output
# targets are mode-agnostic and always requested.

gene_fetch_targets = (
    [os.path.join(main_output_dir, "02_references/sequence_references.csv")]
    if run_gene_fetch else []
)

# ---- PE (merge + concat) targets ----
pe_targets = [
    # Merge mode preprocessing
    os.path.join(preprocessing_dir_merge, "logs/clean_headers/clean_headers.log"),
    # Concat mode preprocessing
    os.path.join(preprocessing_dir_concat, "logs/concat/concat_reads.log"),
    os.path.join(preprocessing_dir_concat, "logs/trim_galore/trim_galore.log"),
    *expand(os.path.join(preprocessing_dir_concat, "trimmed_data/{sample}/{sample}_R1_trimmed.fastq.gz"), sample=list(samples.keys())),
    *expand(os.path.join(preprocessing_dir_concat, "trimmed_data/{sample}/{sample}_R2_trimmed.fastq.gz"), sample=list(samples.keys())),
    # Merge mode barcode recovery
    os.path.join(barcode_recovery_dir_merge, f"consensus/{run_name}_cons_combined-merge.fasta"),
    os.path.join(barcode_recovery_dir_merge, "logs/mge/alignment_files.log"),
    os.path.join(barcode_recovery_dir_merge, f"{run_name}_merge-stats.csv"),
    os.path.join(barcode_recovery_dir_merge, "logs/fasta_cleaner_complete.txt"),
    os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/combined_statistics.csv"),
    os.path.join(barcode_recovery_dir_merge, "logs/exonerate_int_cleanup_complete.txt"),
    os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/01_human_filtered/human_filtered.txt"),
    os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/01_human_filtered/human_filter_metrics.csv"),
    os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/02_at_filtered/at_filtered.txt"),
    os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/02_at_filtered/at_filter_summary.csv"),
    os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/03_outlier_filtered/outlier_filtered.txt"),
    os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/03_outlier_filtered/outlier_filter_summary_metrics.csv"),
    *([os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/04_reference_filtered/reference_filtered.txt"),
       os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/04_reference_filtered/reference_filter_metrics.csv")] if use_reference_filtering else []),
    os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/cleaned_cons_combined.fasta"),
    os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/05_cleaned_consensus/cleaned_cons_metrics-merge.csv"),
    os.path.join(preprocessing_dir_merge, "fastp_summary-merge.csv"),
    # Concat mode barcode recovery
    os.path.join(barcode_recovery_dir_concat, f"consensus/{run_name}_cons_combined-concat.fasta"),
    os.path.join(barcode_recovery_dir_concat, "logs/mge/alignment_files.log"),
    os.path.join(barcode_recovery_dir_concat, f"{run_name}_concat-stats.csv"),
    os.path.join(barcode_recovery_dir_concat, "logs/fasta_cleaner_complete.txt"),
    os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/combined_statistics.csv"),
    os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/01_human_filtered/human_filtered.txt"),
    os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/01_human_filtered/human_filter_metrics.csv"),
    os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/02_at_filtered/at_filtered.txt"),
    os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/02_at_filtered/at_filter_summary.csv"),
    os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/03_outlier_filtered/outlier_filtered.txt"),
    os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/03_outlier_filtered/outlier_filter_summary_metrics.csv"),
    *([os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/04_reference_filtered/reference_filtered.txt"),
       os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/04_reference_filtered/reference_filter_metrics.csv")] if use_reference_filtering else []),
    os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/cleaned_cons_combined.fasta"),
    os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/05_cleaned_consensus/cleaned_cons_metrics-concat.csv"),
    os.path.join(preprocessing_dir_concat, "fastp_summary-concat.csv"),
    os.path.join(barcode_recovery_dir_concat, "logs/exonerate_int_cleanup_complete.txt"),
    # Extra
    os.path.join(barcode_recovery_dir_merge, "logs/fasta_cleaner/fasta_cleaner_complete.txt"),
    os.path.join(barcode_recovery_dir_concat, "logs/fasta_cleaner/fasta_cleaner_complete.txt"),
    os.path.join(barcode_recovery_dir_merge, "logs/gzip_merged_clean_complete.txt"),
]

# ---- SE (Ultima single-end) targets ----
se_targets = [
    os.path.join(preprocessing_dir_se, "fastp_summary-se.csv"),
    os.path.join(barcode_recovery_dir_se, f"consensus/{run_name}_cons_combined-se.fasta"),
    os.path.join(barcode_recovery_dir_se, "logs/mge/alignment_files.log"),
    os.path.join(barcode_recovery_dir_se, f"{run_name}_se-stats.csv"),
    os.path.join(barcode_recovery_dir_se, "logs/fasta_cleaner_complete.txt"),
    os.path.join(barcode_recovery_dir_se, "fasta_cleaner/combined_statistics.csv"),
    os.path.join(barcode_recovery_dir_se, "logs/exonerate_int_cleanup_complete.txt"),
    os.path.join(barcode_recovery_dir_se, "fasta_cleaner/01_human_filtered/human_filtered.txt"),
    os.path.join(barcode_recovery_dir_se, "fasta_cleaner/01_human_filtered/human_filter_metrics.csv"),
    os.path.join(barcode_recovery_dir_se, "fasta_cleaner/02_at_filtered/at_filtered.txt"),
    os.path.join(barcode_recovery_dir_se, "fasta_cleaner/02_at_filtered/at_filter_summary.csv"),
    os.path.join(barcode_recovery_dir_se, "fasta_cleaner/03_outlier_filtered/outlier_filtered.txt"),
    os.path.join(barcode_recovery_dir_se, "fasta_cleaner/03_outlier_filtered/outlier_filter_summary_metrics.csv"),
    *([os.path.join(barcode_recovery_dir_se, "fasta_cleaner/04_reference_filtered/reference_filtered.txt"),
       os.path.join(barcode_recovery_dir_se, "fasta_cleaner/04_reference_filtered/reference_filter_metrics.csv")] if use_reference_filtering else []),
    os.path.join(barcode_recovery_dir_se, "fasta_cleaner/cleaned_cons_combined.fasta"),
    os.path.join(barcode_recovery_dir_se, "fasta_cleaner/05_cleaned_consensus/cleaned_cons_metrics-se.csv"),
    os.path.join(barcode_recovery_dir_se, "logs/fasta_cleaner/fasta_cleaner_complete.txt"),
]

# ---- Mode-agnostic (always requested) targets ----
shared_targets = [
    os.path.join(barcode_recovery_dir, "barcode_consensus_count.tsv"),
    os.path.join(barcode_recovery_dir, "logs/barcode_consensus_count.log"),
    # Structural validation
    os.path.join(main_output_dir, "04_barcode_validation/structural/structural_validation.csv"),
    os.path.join(main_output_dir, "04_barcode_validation/structural/output_barcode_all_passing.fasta"),
    # Taxonomic validation
    os.path.join(main_output_dir, "04_barcode_validation/taxonomic/metrics/01_local_blast_output.csv"),
    os.path.join(main_output_dir, "04_barcode_validation/taxonomic/validated_barcodes.fasta"),
    os.path.join(main_output_dir, "04_barcode_validation/taxonomic/metrics/02_taxonomic_validation.csv"),
    # Reporting
    os.path.join(taxdump_dir, "nodes.dmp"),
    os.path.join(taxdump_dir, "names.dmp"),
    os.path.join(main_output_dir, "05_barcoding_outcome/plots/plots_complete.txt"),
    os.path.join(main_output_dir, "05_barcoding_outcome/multiqc_report/mqc_in_data/mqc_complete.txt"),
    os.path.join(main_output_dir, "05_barcoding_outcome/multiqc_report/multiqc_report.html"),
    # Final outputs
    os.path.join(barcode_recovery_dir, f"{run_name}_barcode_recovery_metrics.csv"),
    os.path.join(barcode_recovery_dir, f"{run_name}_all_cons_combined.fasta"),
    os.path.join(main_output_dir, f"{run_name}_final_metrics.csv"),
    os.path.join(main_output_dir, "05_barcoding_outcome/barcoding_outcome.tsv"),
]

# ---- Final cleanup (PE only; cleanup_files operates on merge+concat dirs) ----
cleanup_targets = [
    os.path.join(preprocessing_dir_concat, "logs/final_cleanup_complete.txt"),
    os.path.join(preprocessing_dir_merge, "logs/final_cleanup_complete.txt"),
]

rule all:
    input:
        gene_fetch_targets,
        (pe_targets if run_mode == "PE" else []),
        (se_targets if run_mode == "SE" else []),
        shared_targets,
        (cleanup_targets if run_mode == "PE" else []),

# ===== GENE FETCH RULE =====
if run_gene_fetch:
    rule gene_fetch:
        input:
            samples_csv=lambda wildcards: config["gene_fetch"].get("samples_file", config["samples_file"])
        output:
            sequence_references=os.path.join(main_output_dir, "02_references/sequence_references.csv"),
            protein_files=expand(
                os.path.join(main_output_dir, "02_references/protein/{sample}.fasta"),
                sample=list(samples.keys())
            )
        params:
            email=config["gene_fetch"]["email"],
            api_key=config["gene_fetch"]["api_key"],
            gene=config["gene_fetch"]["gene"],
            type=config["gene_fetch"]["type"],
            input_type=config["gene_fetch"]["input_type"].lower(),
            genbank=config["gene_fetch"].get("genbank", False),
            output_dir=os.path.join(main_output_dir, "02_references"),
            min_length=config["gene_fetch"].get("minimum_length", False),
            genbank_flag="--genbank" if config["gene_fetch"].get("genbank", False) else "",
            input_flag="--in" if config["gene_fetch"]["input_type"].lower() == "taxid" else "--in2"
        log:
            os.path.join(main_output_dir, "02_references/gene_fetch.log")
        resources:
            mem_mb=lambda wildcards, attempt: rule_resources["gene_fetch"]["mem_mb"] * attempt,
            slurm_partition=rule_resources["gene_fetch"]["partition"]
        retries: 3
        shell:
            """
            set -euo pipefail
            
            echo "Starting gene-fetch: $(date)" > {log}
            echo "Email: {params.email}" >> {log}
            echo "Gene: {params.gene}" >> {log}
            echo "Type: {params.type}" >> {log}
            echo "Input type: {params.input_type}" >> {log}
            echo "Genbank: {params.genbank}" >> {log}
            echo "Input samples: {input.samples_csv}" >> {log}
            echo "Output directory: {params.output_dir}" >> {log}
            
            # Run gene-fetch
            gene-fetch \\
                --email {params.email} \\
                --api-key {params.api_key} \\
                {params.input_flag} {input.samples_csv} \\
                --gene {params.gene} \\
                --type {params.type} \\
                --protein-size {params.min_length} \\
                {params.genbank_flag} \\
                --out {params.output_dir} \\
                >> {log} 2>&1
            
            echo "Gene-fetch completed: $(date)" >> {log}
            echo "Output file: {output.sequence_references}" >> {log}
            
            # Verify output file was created
            if [ ! -f "{output.sequence_references}" ]; then
                echo "ERROR: Expected output file not found: {output.sequence_references}" >> {log}
                exit 1
            fi
            """

# ===== PREPROCESSING RULES - MERGE MODE =====
rule fastp_pe_merge:
    input:
        R1=lambda wildcards: samples[wildcards.sample]["R1"],
        R2=lambda wildcards: samples[wildcards.sample]["R2"]
    output:
        R1_trimmed=temp(os.path.join(preprocessing_dir_merge, "trimmed_data/{sample}/{sample}_R1_trimmed.fastq.gz")),
        R2_trimmed=temp(os.path.join(preprocessing_dir_merge, "trimmed_data/{sample}/{sample}_R2_trimmed.fastq.gz")),
        report=os.path.join(preprocessing_dir_merge, "trimmed_data/{sample}/{sample}_fastp_report.html"),
        json=os.path.join(preprocessing_dir_merge, "trimmed_data/{sample}/{sample}_fastp_report.json"),
        merged_reads=temp(os.path.join(preprocessing_dir_merge, "trimmed_data/{sample}/{sample}_merged.fastq.gz")),
        unpaired_R1=temp(os.path.join(preprocessing_dir_merge, "trimmed_data/unpaired/{sample}/{sample}_unpaired_R1.fastq")),
        unpaired_R2=temp(os.path.join(preprocessing_dir_merge, "trimmed_data/unpaired/{sample}/{sample}_unpaired_R2.fastq"))
    params:
        adapter_r1=fastp_adapter_r1,
        adapter_r2=fastp_adapter_r2,
        extra_args=fastp_extra_args
    log:
        out=os.path.join(preprocessing_dir_merge, "logs/fastp/{sample}_fastp.out"),
        err=os.path.join(preprocessing_dir_merge, "logs/fastp/{sample}_fastp.err")
    threads: rule_resources["fastp_qc"]["threads"]
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["fastp_qc"]["mem_mb"] * attempt
    retries: 3
    shell:
        """
        set -euo pipefail
        
        echo "======== Fastp Merge Processing ========" > {log.out}
        echo "Sample: {wildcards.sample}" >> {log.out}
        echo "Started: $(date)" >> {log.out}
        echo "Input files: {input.R1}, {input.R2}" >> {log.out}
        echo "Threads: {threads}" >> {log.out}
        echo "Adapter R1: {params.adapter_r1}" >> {log.out}
        echo "Adapter R2: {params.adapter_r2}" >> {log.out}
        echo "Extra args: {params.extra_args}" >> {log.out}
        echo "=========================================" >> {log.out}
        
        # Check input files exist
        if [[ ! -f "{input.R1}" ]] || [[ ! -f "{input.R2}" ]]; then
            echo "ERROR: Input files not found" >> {log.out}
            exit 1
        fi
        
        # Log input file sizes
        echo "Input file sizes:" >> {log.out}
        echo "  R1: $(numfmt --to=iec $(stat -c%s {input.R1}))" >> {log.out}
        echo "  R2: $(numfmt --to=iec $(stat -c%s {input.R2}))" >> {log.out}
        
        # Run fastp (merged output written as gzip via .gz extension)
        echo "Running fastp: $(date)" >> {log.out}
        if fastp -i {input.R1} -I {input.R2} \\
              -o {output.R1_trimmed} -O {output.R2_trimmed} \\
              --adapter_sequence {params.adapter_r1} \\
              --adapter_sequence_r2 {params.adapter_r2} \\
              --dedup \\
              --trim_poly_g \\
              --thread {threads} \\
              --merge --merged_out {output.merged_reads} \\
              --unpaired1 {output.unpaired_R1} \\
              --unpaired2 {output.unpaired_R2} \\
              -h {output.report} -j {output.json} \\
              {params.extra_args} \\
              >> {log.out} 2>> {log.err}; then
            echo "Fastp completed successfully: $(date)" >> {log.out}
        else
            echo "ERROR: Fastp failed with exit code $?" >> {log.out}
            exit 1
        fi
        
        # Validate outputs - create mock files if empty or missing
        if [[ ! -s {output.merged_reads} ]]; then
            echo "ERROR: Output files are empty or missing" >> {log.out}
            echo "  Likely cause: Input files were empty, or all reads were filtered out during trimming/merging" >> {log.out}
            echo "Creating empty placeholder outputs..." >> {log.out}
            touch {output.merged_reads}
            touch {output.R1_trimmed}
            touch {output.R2_trimmed}
            echo "Placeholders created: $(date)" >> {log.out}
        else
            echo "Output file sizes:" >> {log.out}
            echo "  R1_trimmed: $(numfmt --to=iec $(stat -c%s {output.R1_trimmed}))" >> {log.out}
            echo "  R2_trimmed: $(numfmt --to=iec $(stat -c%s {output.R2_trimmed}))" >> {log.out}
            echo "  Merged: $(numfmt --to=iec $(stat -c%s {output.merged_reads}))" >> {log.out}
        fi
        echo "Processing completed: $(date)" >> {log.out}
        """

if downsampling_enabled:
    rule downsample_merge:
        input:
            merged_reads=os.path.join(preprocessing_dir_merge, "trimmed_data/{sample}/{sample}_merged.fastq.gz")
        output:
            downsampled=temp(os.path.join(preprocessing_dir_merge, "trimmed_data/{sample}/{sample}_merged_downsampled.fastq.gz"))
        params:
            max_reads=config["downsampling"]["max_reads"]
        log:
            out=os.path.join(preprocessing_dir_merge, "logs/downsample/{sample}.out"),
            err=os.path.join(preprocessing_dir_merge, "logs/downsample/{sample}.err")
        resources:
            mem_mb=lambda wildcards, attempt: rule_resources["downsample"]["mem_mb"] * attempt
        retries: 3
        shell:
            """
            set -euo pipefail
        
            echo "Downsampling merged reads for {wildcards.sample}" > {log.out}
            echo "Started: $(date)" >> {log.out}
            echo "Target reads: {params.max_reads}" >> {log.out}
            
            # Check if input file is empty or missing
            if [ ! -s "{input.merged_reads}" ]; then
                echo "Input file is empty or missing - creating placeholder output" >> {log.out}
                touch {output.downsampled}
                echo "Placeholder output created: $(date)" >> {log.out}
                exit 0
            fi
            
            reformat.sh -Xmx$(( {resources.mem_mb} * 80 / 100 ))m -eoom \\
                        in={input.merged_reads} \\
                        out={output.downsampled} \\
                        samplereadstarget={params.max_reads} \\
                        sampleseed=12345 \\
                        >> {log.out} 2>> {log.err}
        
            echo "Completed: $(date)" >> {log.out}
            echo "Output file size: $(numfmt --to=iec $(stat -c%s {output.downsampled}))" >> {log.out}
            """

rule clean_headers_merge:
    input:
        merged_reads=os.path.join(preprocessing_dir_merge, "trimmed_data/{sample}/{sample}_merged_downsampled.fastq.gz") if downsampling_enabled else os.path.join(preprocessing_dir_merge, "trimmed_data/{sample}/{sample}_merged.fastq.gz")
    output:
        clean_merged=os.path.join(preprocessing_dir_merge, "trimmed_data/{sample}/{sample}_merged_clean.fastq")
    log:
        out=temp(os.path.join(preprocessing_dir_merge, "logs/clean_headers/{sample}.out")),
        err=temp(os.path.join(preprocessing_dir_merge, "logs/clean_headers/{sample}.err"))
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["clean_headers_merge"]["mem_mb"] * attempt
    retries: 3
    shell:
        """
        set -euo pipefail
        
        # Create directories
        mkdir -p $(dirname {log.out})
        mkdir -p $(dirname {output.clean_merged})
        
        # Create log header
        echo "Cleaning headers for {wildcards.sample}" > {log.out}
        echo "Started: $(date)" >> {log.out}
        echo "Host: $(hostname)" >> {log.out}
        echo "Input: {input.merged_reads}" >> {log.out}
        
        # Check if input file is empty or missing
        if [ ! -s "{input.merged_reads}" ]; then
            echo "Input file is empty or missing - creating placeholder output" >> {log.out}
            touch {output.clean_merged}
            echo "Placeholder output created: $(date)" >> {log.out}
        else
            # Decompress and clean headers in one step
            if zcat {input.merged_reads} | sed 's/ /_/g' > {output.clean_merged} 2>> {log.err}; then
                echo "Completed: $(date)" >> {log.out}
                echo "Output file size: $(numfmt --to=iec $(stat -c%s {output.clean_merged}))" >> {log.out}
                echo "--------------------------------------------" >> {log.out}
            else
                echo "FAILED: $(date)" >> {log.out}
                exit 1
            fi
        fi
        """

rule aggregate_clean_headers_logs_merge:
    input:
        sample_logs=expand(os.path.join(preprocessing_dir_merge, "logs/clean_headers/{sample}.out"), sample=list(samples.keys()))
    output:
        combined_log=os.path.join(preprocessing_dir_merge, "logs/clean_headers/clean_headers.log") 
    log:
        out=os.path.join(preprocessing_dir_merge, "logs/clean_headers/aggregate_clean_headers.out"),
        err=os.path.join(preprocessing_dir_merge, "logs/clean_headers/aggregate_clean_headers.err")
    threads: 1
    resources:
        mem_mb=3076
    retries: 3
    shell:
        """
        set -euo pipefail
        
        # Create header for aggregated log file
        echo "===============================================" > {output.combined_log}
        echo "Aggregated clean headers logs - Created $(date)" >> {output.combined_log}
        echo "===============================================" >> {output.combined_log}
        echo "" >> {output.combined_log}
        
        # Concatenate all sample logs into a combined log
        cat {input.sample_logs} >> {output.combined_log} 2>> {log.err}
        echo "Aggregation complete." >> {log.out}
        """

# ===== PREPROCESSING RULES - CONCAT MODE =====
rule fastp_pe_concat:
    input:
        R1=lambda wildcards: samples[wildcards.sample]["R1"],
        R2=lambda wildcards: samples[wildcards.sample]["R2"]
    output:
        R1_trimmed=os.path.join(preprocessing_dir_concat, "trimmed_data/{sample}/{sample}_R1_trimmed.fastq.gz"),
        R2_trimmed=os.path.join(preprocessing_dir_concat, "trimmed_data/{sample}/{sample}_R2_trimmed.fastq.gz"),
        report=os.path.join(preprocessing_dir_concat, "trimmed_data/{sample}/{sample}_fastp_report.html"),
        json=os.path.join(preprocessing_dir_concat, "trimmed_data/{sample}/{sample}_fastp_report.json")
    params:
        adapter_r1=fastp_adapter_r1,
        adapter_r2=fastp_adapter_r2,
        extra_args=fastp_extra_args
    log:
        out=os.path.join(preprocessing_dir_concat, "logs/fastp/{sample}_fastp.out"),
        err=os.path.join(preprocessing_dir_concat, "logs/fastp/{sample}_fastp.err")
    threads: rule_resources["fastp_qc"]["threads"]
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["fastp_qc"]["mem_mb"] * attempt
    retries: 3
    shell:
        """
        set -euo pipefail
        
        # Create output directories
        mkdir -p $(dirname {output.R1_trimmed})
        mkdir -p $(dirname {output.report})
        
        # Log job start
        echo "======== Fastp Concat Processing ========" > {log.out}
        echo "Sample: {wildcards.sample}" >> {log.out}
        echo "Started: $(date)" >> {log.out}
        echo "Input files: {input.R1}, {input.R2}" >> {log.out}
        echo "Threads: {threads}" >> {log.out}
        echo "Adapter R1: {params.adapter_r1}" >> {log.out}
        echo "Adapter R2: {params.adapter_r2}" >> {log.out}
        echo "Extra args: {params.extra_args}" >> {log.out}
        echo "=========================================" >> {log.out}
        
        # Check input files exist
        if [[ ! -f "{input.R1}" ]] || [[ ! -f "{input.R2}" ]]; then
            echo "ERROR: Input files not found" >> {log.out}
            exit 1
        fi
        
        # Log input file sizes
        echo "Input file sizes:" >> {log.out}
        echo "  R1: $(numfmt --to=iec $(stat -c%s {input.R1}))" >> {log.out}
        echo "  R2: $(numfmt --to=iec $(stat -c%s {input.R2}))" >> {log.out}
        
        # Run fastp (outputs gzipped natively via .gz extension auto-detection)
        echo "Running fastp: $(date)" >> {log.out}
        if fastp -i {input.R1} -I {input.R2} \\
                 -o {output.R1_trimmed} -O {output.R2_trimmed} \\
                 --adapter_sequence {params.adapter_r1} \\
                 --adapter_sequence_r2 {params.adapter_r2} \\
                 --dedup \\
                 --trim_poly_g \\
                 --thread {threads} \\
                 -h {output.report} -j {output.json} \\
                 {params.extra_args} \\
                 >> {log.out} 2>> {log.err}; then
            echo "Fastp completed successfully: $(date)" >> {log.out}
        else
            echo "ERROR: Fastp failed with exit code $?" >> {log.out}
            exit 1
        fi
        
        # Validate outputs - create mock files if empty or missing
        if [[ ! -s {output.R1_trimmed} ]] || [[ ! -s {output.R2_trimmed} ]]; then
            echo "ERROR: Output files are empty or missing" >> {log.out}
            echo "  Likely cause: Input files were empty, or all reads were filtered out during trimming" >> {log.out}
            echo "Creating empty placeholder outputs..." >> {log.out}
            touch {output.R1_trimmed}
            touch {output.R2_trimmed}
            echo "Placeholders created: $(date)" >> {log.out}
        else
            echo "Output file sizes:" >> {log.out}
            echo "  R1_trimmed: $(numfmt --to=iec $(stat -c%s {output.R1_trimmed}))" >> {log.out}
            echo "  R2_trimmed: $(numfmt --to=iec $(stat -c%s {output.R2_trimmed}))" >> {log.out}
        fi
        
        echo "Processing completed: $(date)" >> {log.out}
        """

rule fastq_concat:
    input:
        R1=os.path.join(preprocessing_dir_concat, "trimmed_data/{sample}/{sample}_R1_trimmed.fastq.gz"),
        R2=os.path.join(preprocessing_dir_concat, "trimmed_data/{sample}/{sample}_R2_trimmed.fastq.gz")
    output:
        temp(os.path.join(preprocessing_dir_concat, "trimmed_data/{sample}/{sample}_concat.fastq"))
    log:
        out=os.path.join(preprocessing_dir_concat, "logs/concat/{sample}.out")
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["fastq_concat"]["mem_mb"] * attempt
    retries: 3
    shell:
        """
        set -euo pipefail
                
        echo "Concatenating and cleaning headers for {wildcards.sample}" > {log.out}
        echo "Started: $(date)" >> {log.out}
        echo "Input R1: {input.R1}" >> {log.out}
        echo "Input R2: {input.R2}" >> {log.out}
        
        # Check if input files are empty or missing
        if [ ! -s "{input.R1}" ] || [ ! -s "{input.R2}" ]; then
            echo "Input files are empty or missing - creating placeholder output" >> {log.out}
            touch {output}
            echo "Placeholder output created: $(date)" >> {log.out}
        else
            # Decompress, clean headers (replace spaces with underscores), and concatenate in one step
            (zcat {input.R1} | sed 's/ /_/g' && zcat {input.R2} | sed 's/ /_/g') > {output} 2>> {log.out}
            echo "Completed: $(date)" >> {log.out}
            echo "Output file size: $(numfmt --to=iec $(stat -c%s {output}))" >> {log.out}
        fi
        """

rule aggregate_concat_logs:
    input:
        sample_logs=expand(os.path.join(preprocessing_dir_concat, "logs/concat/{sample}.out"), sample=list(samples.keys()))
    output:
        combined_log=os.path.join(preprocessing_dir_concat, "logs/concat/concat_reads.log")
    log:
        out=os.path.join(preprocessing_dir_concat, "logs/concat/aggregate_concat.out"),
        err=os.path.join(preprocessing_dir_concat, "logs/concat/aggregate_concat.err")
    threads: 1
    resources:
        mem_mb=3076
    retries: 3
    shell:
        """
        set -euo pipefail

        echo "===============================================" > {output.combined_log}
        echo "Aggregated concat logs - Created $(date)" >> {output.combined_log}
        echo "===============================================" >> {output.combined_log}
        echo "" >> {output.combined_log}

        # Concatenate all sample logs into a combined log
        cat {input.sample_logs} >> {output.combined_log} 2>> {log.err}
        echo "Aggregation complete." >> {log.out}
        """

rule quality_trim:
    input:
        os.path.join(preprocessing_dir_concat, "trimmed_data/{sample}/{sample}_concat.fastq")
    output:
        concat_trimmed=temp(os.path.join(preprocessing_dir_concat, "trimmed_data/{sample}/{sample}_concat_trimmed.fq")),
        report=os.path.join(preprocessing_dir_concat, "trimmed_data/{sample}/{sample}_concat.fastq_trimming_report.txt")
    params:
        outdir=os.path.join(preprocessing_dir_concat, "trimmed_data/{sample}"),
        report_dir=os.path.join(preprocessing_dir_concat, "trimmed_data/")
    log:
        out=os.path.join(preprocessing_dir_concat, "logs/trim_galore/{sample}.out"),
        err=os.path.join(preprocessing_dir_concat, "logs/trim_galore/{sample}.err")
    threads: rule_resources["quality_trim"]["threads"]
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["quality_trim"]["mem_mb"] * attempt
    retries: 3
    shell:
        """
        set -euo pipefail

        echo "Quality trimming for {wildcards.sample}" > {log.out}
        echo "Started: $(date)" >> {log.out}
        echo "Input file: {input}" >> {log.out}

        # Check if input file is empty or missing
        if [ ! -s "{input}" ]; then
            echo "Input file is empty or missing - creating placeholder outputs" >> {log.out}
            touch {output.concat_trimmed}
            touch {output.report}
            echo "Placeholder outputs created: $(date)" >> {log.out}
            exit 0
        fi

        # Run trim_galore
        trim_galore --cores {threads} \\
                    --dont_gzip \\
                    --output_dir {params.outdir} \\
                    --basename {wildcards.sample}_concat \\
                    {input} >> {log.out} 2>> {log.err}

        # Log job completion
        echo "Completed: $(date)" >> {log.out}
        echo "Output file size: $(numfmt --to=iec $(stat -c%s {output[0]}))" >> {log.out}
        """

if downsampling_enabled:
    rule downsample_concat:
        input:
            concat_trimmed=os.path.join(preprocessing_dir_concat, "trimmed_data/{sample}/{sample}_concat_trimmed.fq")
        output:
            downsampled=temp(os.path.join(preprocessing_dir_concat, "trimmed_data/{sample}/{sample}_concat_trimmed_downsampled.fastq"))
        params:
            max_reads=config["downsampling"]["max_reads"] * 2  # Multiply by 2 because concat mode has R1+R2 reads
        log:
            out=os.path.join(preprocessing_dir_concat, "logs/downsample/{sample}.out"),
            err=os.path.join(preprocessing_dir_concat, "logs/downsample/{sample}.err")
        resources:
            mem_mb=lambda wildcards, attempt: rule_resources["downsample"]["mem_mb"] * attempt
        retries: 3
        shell:
            """
            set -euo pipefail

            echo "Downsampling concatenated trimmed reads for {wildcards.sample}" > {log.out}
            echo "Started: $(date)" >> {log.out}
            echo "Target reads: {params.max_reads} (2x read pairs for concat mode)" >> {log.out}
            
            # Check if input file is empty or missing
            if [ ! -s "{input.concat_trimmed}" ]; then
                echo "Input file is empty or missing - creating placeholder output" >> {log.out}
                touch {output.downsampled}
                echo "Placeholder output created: $(date)" >> {log.out}
                exit 0
            fi

            reformat.sh -Xmx$(( {resources.mem_mb} * 80 / 100 ))m -eoom \\
                        in={input.concat_trimmed} \\
                        out={output.downsampled} \\
                        samplereadstarget={params.max_reads} \\
                        sampleseed=12345 \\
                        >> {log.out} 2>> {log.err}

            echo "Completed: $(date)" >> {log.out}
            echo "Output file size: $(numfmt --to=iec $(stat -c%s {output.downsampled}))" >> {log.out}
            """

rule aggregate_trim_galore_logs:
    input:
        sample_logs=expand(os.path.join(preprocessing_dir_concat, "logs/trim_galore/{sample}.out"), sample=list(samples.keys()))
    output:
        combined_log=os.path.join(preprocessing_dir_concat, "logs/trim_galore/trim_galore.log")
    log:
        out=os.path.join(preprocessing_dir_concat, "logs/aggregate_trim_galore.out"),
        err=os.path.join(preprocessing_dir_concat, "logs/aggregate_trim_galore.err")
    threads: 1
    resources:
        mem_mb=3076
    shell:
        """
        set -euo pipefail

        echo "===============================================" > {output.combined_log}
        echo "Aggregated trim_galore logs - Created $(date)" >> {output.combined_log}
        echo "===============================================" >> {output.combined_log}
        echo "" >> {output.combined_log}

        # Concatenate all sample logs into a combined log
        cat {input.sample_logs} >> {output.combined_log} 2>> {log.err}
        echo "Aggregation complete." >> {log.out}
        """


# ===== PREPROCESSING RULES - SE (ULTIMA) MODE =====
rule fastp_se:
    input:
        R1=lambda wildcards: samples[wildcards.sample]["R1"],
        overrep_fasta=os.path.join(workflow.basedir, "..", "resources/overrep_fasta/overrep.fasta")        # Depreciated. Left in case of future application
    output:
        trimmed=os.path.join(preprocessing_dir_se, "trimmed_data/{sample}/{sample}_se_trimmed.fastq"),
        report=os.path.join(preprocessing_dir_se, "trimmed_data/{sample}/{sample}_fastp_report.html"),
        json=os.path.join(preprocessing_dir_se, "trimmed_data/{sample}/{sample}_fastp_report.json")
    params:
        extra_args=fastp_extra_args,
        raw=os.path.join(preprocessing_dir_se, "trimmed_data/{sample}/{sample}_se_trimmed.raw.fastq")
    log:
        out=os.path.join(preprocessing_dir_se, "logs/fastp/{sample}_fastp.out"),
        err=os.path.join(preprocessing_dir_se, "logs/fastp/{sample}_fastp.err")
    threads: rule_resources["fastp_qc"]["threads"]
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["fastp_qc"]["mem_mb"] * attempt
    retries: 3
    shell:
        """
        set -euo pipefail

        mkdir -p $(dirname {output.trimmed})
        mkdir -p $(dirname {log.out})

        echo "======== Fastp SE (Ultima) Processing ========" > {log.out}
        echo "Sample: {wildcards.sample}" >> {log.out}
        echo "Started: $(date)" >> {log.out}
        echo "Input file: {input.R1}" >> {log.out}
        echo "Threads: {threads}" >> {log.out}
        echo "Adapter detection: auto (Ultima SE; -a omitted)" >> {log.out}
        echo "Poly-X trimming: ENABLED (--trim_poly_x)" >> {log.out}
        echo "Poly-G trimming: DISABLED (-G; Ultima is not two-colour)" >> {log.out}
        echo "Extra args: {params.extra_args}" >> {log.out}
        echo "==============================================" >> {log.out}

        if [[ ! -f "{input.R1}" ]]; then
            echo "ERROR: Input file not found: {input.R1}" >> {log.out}
            exit 1
        fi

        echo "Input file size: $(numfmt --to=iec $(stat -c%s {input.R1}))" >> {log.out}

        # Run fastp (single-end). Poly-G trimming explicitly disabled with -G.
        echo "Running fastp: $(date)" >> {log.out}
        if fastp -i {input.R1} \\
                 -o {params.raw} \\
                 --dedup \\
                 --trim_poly_x \\
                 --poly_x_min_len 6 \\
                 -p  \\
                 -G \\
                 --thread {threads} \\
                 -h {output.report} -j {output.json} \\
                 {params.extra_args} \\
                 >> {log.out} 2>> {log.err}; then
            echo "Fastp completed successfully: $(date)" >> {log.out}
        else
            echo "ERROR: Fastp failed with exit code $?" >> {log.out}
            exit 1
        fi

        # Clean read headers (replace spaces with underscores) for MGE compatibility.
        if [[ -s {params.raw} ]]; then
            sed 's/ /_/g' {params.raw} > {output.trimmed}
            rm -f {params.raw}
            echo "Trimmed output size: $(numfmt --to=iec $(stat -c%s {output.trimmed}))" >> {log.out}
        else
            echo "WARNING: fastp produced empty output (all reads filtered or empty input)" >> {log.out}
            echo "Creating empty placeholder output..." >> {log.out}
            touch {output.trimmed}
            rm -f {params.raw}
        fi

        echo "Processing completed: $(date)" >> {log.out}
        """

# ----- SE downsampling (optional; mirrors downsample_concat) -----
if downsampling_enabled:
    rule downsample_se:
        input:
            se_trimmed=os.path.join(preprocessing_dir_se, "trimmed_data/{sample}/{sample}_se_trimmed.fastq")
        output:
            downsampled=temp(os.path.join(preprocessing_dir_se, "trimmed_data/{sample}/{sample}_se_trimmed_downsampled.fastq"))
        params:
            max_reads=config["downsampling"]["max_reads"]  # single-end: reads, not pairs
        log:
            out=os.path.join(preprocessing_dir_se, "logs/downsample/{sample}.out"),
            err=os.path.join(preprocessing_dir_se, "logs/downsample/{sample}.err")
        resources:
            mem_mb=lambda wildcards, attempt: rule_resources["downsample"]["mem_mb"] * attempt
        retries: 3
        shell:
            """
            set -euo pipefail

            mkdir -p $(dirname {log.out})
            echo "Downsampling SE trimmed reads for {wildcards.sample}" > {log.out}
            echo "Started: $(date)" >> {log.out}
            echo "Target reads: {params.max_reads}" >> {log.out}

            if [ ! -s "{input.se_trimmed}" ]; then
                echo "Input file is empty or missing - creating placeholder output" >> {log.out}
                touch {output.downsampled}
                echo "Placeholder output created: $(date)" >> {log.out}
                exit 0
            fi

            reformat.sh -Xmx$(( {resources.mem_mb} * 80 / 100 ))m -eoom \\
                        in={input.se_trimmed} \\
                        out={output.downsampled} \\
                        samplereadstarget={params.max_reads} \\
                        sampleseed=12345 \\
                        >> {log.out} 2>> {log.err}

            echo "Completed: $(date)" >> {log.out}
            echo "Output file size: $(numfmt --to=iec $(stat -c%s {output.downsampled}))" >> {log.out}
            """

# ----- SE fastp summary (mirrors fastp_summary_concat) -----
rule fastp_summary_se:
    input:
        json_files=expand(
            os.path.join(preprocessing_dir_se, "trimmed_data/{sample}/{sample}_fastp_report.json"),
            sample=list(samples.keys())
        )
    output:
        summary_csv=os.path.join(preprocessing_dir_se, "fastp_summary-se.csv")
    params:
        script=os.path.join(workflow.basedir, "scripts/fastp_summary_parser.py"),
        input_dir=os.path.join(preprocessing_dir_se, "trimmed_data")
    log:
        os.path.join(preprocessing_dir_se, "logs/fastp_summary_se.log")
    threads: 1
    resources:
        mem_mb=4096
    retries: 2
    shell:
        """
        set -euo pipefail

        mkdir -p $(dirname {log})

        echo "Generating fastp summary CSV for SE mode: $(date)" > {log}
        echo "Input directory: {params.input_dir}" >> {log}
        echo "Output CSV: {output.summary_csv}" >> {log}

        python {params.script} \\
            -i {params.input_dir} \\
            -o {output.summary_csv} \\
            >> {log} 2>&1

        echo "Completed: $(date)" >> {log}
        """


# ===== MGE RULES FOR BARCODE RECOVERY =====

# MGE rule for merge mode
rule MitoGeneExtractor_merge:
    input:
        DNA=get_mge_input_merge,
        AA=get_protein_reference
    output:
        alignment=os.path.join(barcode_recovery_dir_merge, "alignment/{sample}_r_{r}_s_{s}_align_{sample}.fas"),
        consensus=os.path.join(barcode_recovery_dir_merge, "consensus/{sample}_r_{r}_s_{s}_con_{sample}.fas")
    log:
        out=os.path.join(barcode_recovery_dir_merge, "out/{sample}_r_{r}_s_{s}_summary.out"),
        err=os.path.join(barcode_recovery_dir_merge, "err/{sample}_r_{r}_s_{s}_summary.err")
    params:
        mge_executor=MGE_EXECUTABLE,
        n=n,
        C=C,
        t=t,
        output_dir=barcode_recovery_dir_merge,
        vulgar_dir=lambda wildcards: os.path.join(barcode_recovery_dir_merge, f"logs/mge/{wildcards.sample}_r_{wildcards.r}_s_{wildcards.s}/")
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["MitoGeneExtractor"]["mem_mb"] * attempt,
        slurm_partition=rule_resources["MitoGeneExtractor"]["partition"]
    retries: 4
    shell:
        """
        set -euo pipefail
        
        # Check if input file is empty or missing
        if [ ! -s "{input.DNA}" ]; then
            echo "===== EMPTY/MISSING INPUT DETECTED =====" >> {log.out}
            echo "Started: $(date)" >> {log.out}
            echo "Sample: {wildcards.sample}" >> {log.out}
            echo "Mode: merge" >> {log.out}
            echo "Input file: {input.DNA}" >> {log.out}
            
            if [ ! -f "{input.DNA}" ]; then
                echo "Input file does not exist" >> {log.out}
            else
                FILE_SIZE=$(stat -c%s "{input.DNA}" 2>/dev/null || echo "unknown")
                echo "Input file size: $FILE_SIZE bytes" >> {log.out}
                echo "File is empty - likely all reads were filtered out during trimming" >> {log.out}
            fi
            
            echo "Creating mock outputs..." >> {log.out}
            echo ">Consensus__{wildcards.sample}" > {output.consensus}
            touch {output.alignment}
            touch {log.err}
            echo "Mock outputs created successfully: $(date)" >> {log.out}
            echo "=========================================" >> {log.out}
            exit 0
        fi
        
        # Create vulgar directory if it doesn't exist
        mkdir -p {params.vulgar_dir}
        
        # Job info and file reporting
        echo "===== Job Info =====" >> {log.out}
        echo "Started: $(date)" >> {log.out}
        echo "Node: $(hostname)" >> {log.out}
        echo "Memory allocated: $(({resources.mem_mb} / 1024))GB" >> {log.out}
        echo "Mode: merge" >> {log.out}
        echo "Sample: {wildcards.sample}" >> {log.out}
        
        # File size analysis
        FILE_SIZE_BYTES=$(stat -c%s "{input.DNA}")
        FILE_SIZE_GB=$(echo "scale=2; $FILE_SIZE_BYTES / 1024 / 1024 / 1024" | bc)
        NUM_SEQUENCES=$(grep -c "^@" "{input.DNA}" || echo "0")
        
        echo "===== Input Files =====" >> {log.out}
        echo "Input file: {input.DNA}" >> {log.out}
        echo "File size: ${{FILE_SIZE_GB}}GB ($FILE_SIZE_BYTES bytes)" >> {log.out}
        echo "Number of sequences: $NUM_SEQUENCES" >> {log.out}        
        echo "Input AA file: {input.AA}" >> {log.out}
            
        # Run MitoGeneExtractor
        echo "===== Starting MitoGeneExtractor =====" >> {log.out}
        echo "Command: {params.mge_executor} -q {input.DNA} -p {input.AA} -o {params.output_dir}/alignment/{wildcards.sample}_r_{wildcards.r}_s_{wildcards.s}_align_ -c {params.output_dir}/consensus/{wildcards.sample}_r_{wildcards.r}_s_{wildcards.s}_con_ -V {params.vulgar_dir} -r {wildcards.r} -s {wildcards.s} -n {params.n} -C {params.C} -t {params.t} --temporaryDirectory {params.output_dir} --verbosity 10" >> {log.out}
        echo "=======================================" >> {log.out}
        
        time {params.mge_executor} \\
        -q {input.DNA} -p {input.AA} \\
        -o {params.output_dir}/alignment/{wildcards.sample}_r_{wildcards.r}_s_{wildcards.s}_align_ \\
        -c {params.output_dir}/consensus/{wildcards.sample}_r_{wildcards.r}_s_{wildcards.s}_con_ \\
        -V {params.vulgar_dir} \\
        -r {wildcards.r} -s {wildcards.s} \\
        -n {params.n} \\
        -C {params.C} \\
        -t {params.t} \\
        --temporaryDirectory {params.output_dir} \\
        --verbosity 10 \\
        >> {log.out} 2>> {log.err} || {{
            EXIT_CODE=$?
            echo "===== MitoGeneExtractor Failed =====" >> {log.out}
            echo "Exit code: $EXIT_CODE" >> {log.out}
            echo "Finished: $(date)" >> {log.out}
            
            # Enhanced failure analysis
            if [ $EXIT_CODE -eq 137 ]; then
                echo "ERROR: Process killed (likely out of memory)" >> {log.out}
                echo "  Current allocation: $(({resources.mem_mb} / 1024))GB" >> {log.out}
                echo "  File size: ${{FILE_SIZE_GB}}GB" >> {log.out}
                echo "  Sequences processed: $NUM_SEQUENCES" >> {log.out}
            elif [ $EXIT_CODE -eq 124 ]; then
                echo "ERROR: Process timed out" >> {log.out}
            elif [ $EXIT_CODE -eq 1 ]; then
                echo "ERROR: General error - check error log for details" >> {log.out}
            else
                echo "ERROR: Unknown failure (exit code $EXIT_CODE)" >> {log.out}
            fi
            
            echo "Final system state:" >> {log.out}
            free -h >> {log.out}
            echo "===================================" >> {log.out}
            
            exit $EXIT_CODE
        }}
        
        echo "===== Job Completed Successfully =====" >> {log.out}
        echo "Finished: $(date)" >> {log.out}
        echo "Final memory usage:" >> {log.out}
        free -h >> {log.out}
        
        # Validate and report output files
        if [ -f "{output.consensus}" ] && [ -f "{output.alignment}" ]; then
            CONSENSUS_SIZE=$(stat -c%s {output.consensus})
            ALIGNMENT_SIZE=$(stat -c%s {output.alignment})
            CONSENSUS_LINES=$(wc -l < {output.consensus})
            
            echo "===== Output Summary =====" >> {log.out}
            echo "Consensus file: $(numfmt --to=iec $CONSENSUS_SIZE) ($CONSENSUS_LINES lines)" >> {log.out}
            echo "Alignment file: $(numfmt --to=iec $ALIGNMENT_SIZE)" >> {log.out}
            
            if [ $CONSENSUS_LINES -gt 1 ]; then
                echo "✓ Consensus contains sequence data" >> {log.out}
            else
                echo "⚠ WARNING: Consensus may be header-only" >> {log.out}
            fi
            echo "===========================" >> {log.out}
        else
            echo "ERROR: Expected output files not found" >> {log.out}
            echo "Expected files:" >> {log.out}
            echo "  Consensus: {output.consensus}" >> {log.out}
            echo "  Alignment: {output.alignment}" >> {log.out}
            exit 1
        fi
        """


# MGE rule for concat mode
rule MitoGeneExtractor_concat:
    input:
        DNA=get_mge_input_concat,
        AA=get_protein_reference
    output:
        alignment=os.path.join(barcode_recovery_dir_concat, "alignment/{sample}_r_{r}_s_{s}_align_{sample}.fas"),
        consensus=os.path.join(barcode_recovery_dir_concat, "consensus/{sample}_r_{r}_s_{s}_con_{sample}.fas")
    log:
        out=os.path.join(barcode_recovery_dir_concat, "out/{sample}_r_{r}_s_{s}_summary.out"),
        err=os.path.join(barcode_recovery_dir_concat, "err/{sample}_r_{r}_s_{s}_summary.err")
    params:
        mge_executor=MGE_EXECUTABLE,
        n=n,
        C=C,
        t=t,
        output_dir=barcode_recovery_dir_concat,
        vulgar_dir=lambda wildcards: os.path.join(barcode_recovery_dir_concat, f"logs/mge/{wildcards.sample}_r_{wildcards.r}_s_{wildcards.s}/")
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["MitoGeneExtractor"]["mem_mb"] * attempt,
        slurm_partition=rule_resources["MitoGeneExtractor"]["partition"]
    retries: 4
    shell:
        """
        set -euo pipefail
        
        # Check if input file is empty or missing
        if [ ! -s "{input.DNA}" ]; then
            echo "===== EMPTY/MISSING INPUT DETECTED =====" >> {log.out}
            echo "Started: $(date)" >> {log.out}
            echo "Sample: {wildcards.sample}" >> {log.out}
            echo "Mode: concat" >> {log.out}
            echo "Input file: {input.DNA}" >> {log.out}
            
            if [ ! -f "{input.DNA}" ]; then
                echo "Input file does not exist" >> {log.out}
            else
                FILE_SIZE=$(stat -c%s "{input.DNA}" 2>/dev/null || echo "unknown")
                echo "Input file size: $FILE_SIZE bytes" >> {log.out}
                echo "File is empty - likely all reads were filtered out during trimming" >> {log.out}
            fi
            
            echo "Creating mock outputs..." >> {log.out}
            echo ">Consensus__{wildcards.sample}" > {output.consensus}
            touch {output.alignment}
            touch {log.err}
            echo "Mock outputs created successfully: $(date)" >> {log.out}
            echo "=========================================" >> {log.out}
            exit 0
        fi
        
        # Create vulgar directory if it doesn't exist
        mkdir -p {params.vulgar_dir}
        
        # Job info and file reporting
        echo "===== Job Info =====" >> {log.out}
        echo "Started: $(date)" >> {log.out}
        echo "Node: $(hostname)" >> {log.out}
        echo "Memory allocated: $(({resources.mem_mb} / 1024))GB" >> {log.out}
        echo "Mode: concat" >> {log.out}
        echo "Sample: {wildcards.sample}" >> {log.out}
        
        # File size analysis
        FILE_SIZE_BYTES=$(stat -c%s "{input.DNA}")
        FILE_SIZE_GB=$(echo "scale=2; $FILE_SIZE_BYTES / 1024 / 1024 / 1024" | bc)
        NUM_SEQUENCES=$(grep -c "^@" "{input.DNA}" || echo "0")
        
        echo "===== Input Files =====" >> {log.out}
        echo "Input file: {input.DNA}" >> {log.out}
        echo "File size: ${{FILE_SIZE_GB}}GB ($FILE_SIZE_BYTES bytes)" >> {log.out}
        echo "Number of sequences: $NUM_SEQUENCES" >> {log.out}        
        echo "Input AA file: {input.AA}" >> {log.out}
            
        # Run MitoGeneExtractor
        echo "===== Starting MitoGeneExtractor =====" >> {log.out}
        echo "Command: {params.mge_executor} -q {input.DNA} -p {input.AA} -o {params.output_dir}/alignment/{wildcards.sample}_r_{wildcards.r}_s_{wildcards.s}_align_ -c {params.output_dir}/consensus/{wildcards.sample}_r_{wildcards.r}_s_{wildcards.s}_con_ -V {params.vulgar_dir} -r {wildcards.r} -s {wildcards.s} -n {params.n} -C {params.C} -t {params.t} --temporaryDirectory {params.output_dir} --verbosity 10" >> {log.out}
        echo "=======================================" >> {log.out}
        
        time {params.mge_executor} \\
        -q {input.DNA} -p {input.AA} \\
        -o {params.output_dir}/alignment/{wildcards.sample}_r_{wildcards.r}_s_{wildcards.s}_align_ \\
        -c {params.output_dir}/consensus/{wildcards.sample}_r_{wildcards.r}_s_{wildcards.s}_con_ \\
        -V {params.vulgar_dir} \\
        -r {wildcards.r} -s {wildcards.s} \\
        -n {params.n} \\
        -C {params.C} \\
        -t {params.t} \\
        --temporaryDirectory {params.output_dir} \\
        --verbosity 10 \\
        >> {log.out} 2>> {log.err} || {{
            EXIT_CODE=$?
            echo "===== MitoGeneExtractor Failed =====" >> {log.out}
            echo "Exit code: $EXIT_CODE" >> {log.out}
            echo "Finished: $(date)" >> {log.out}
            
            # Enhanced failure analysis
            if [ $EXIT_CODE -eq 137 ]; then
                echo "ERROR: Process killed (likely out of memory)" >> {log.out}
                echo "  Current allocation: $(({resources.mem_mb} / 1024))GB" >> {log.out}
                echo "  File size: ${{FILE_SIZE_GB}}GB" >> {log.out}
                echo "  Sequences processed: $NUM_SEQUENCES" >> {log.out}
            elif [ $EXIT_CODE -eq 124 ]; then
                echo "ERROR: Process timed out" >> {log.out}
            elif [ $EXIT_CODE -eq 1 ]; then
                echo "ERROR: General error - check error log for details" >> {log.out}
            else
                echo "ERROR: Unknown failure (exit code $EXIT_CODE)" >> {log.out}
            fi
            
            echo "Final system state:" >> {log.out}
            free -h >> {log.out}
            echo "===================================" >> {log.out}
            
            exit $EXIT_CODE
        }}
        
        echo "===== Job Completed Successfully =====" >> {log.out}
        echo "Finished: $(date)" >> {log.out}
        echo "Final memory usage:" >> {log.out}
        free -h >> {log.out}
        
        # Validate and report output files
        if [ -f "{output.consensus}" ] && [ -f "{output.alignment}" ]; then
            CONSENSUS_SIZE=$(stat -c%s {output.consensus})
            ALIGNMENT_SIZE=$(stat -c%s {output.alignment})
            CONSENSUS_LINES=$(wc -l < {output.consensus})
            
            echo "===== Output Summary =====" >> {log.out}
            echo "Consensus file: $(numfmt --to=iec $CONSENSUS_SIZE) ($CONSENSUS_LINES lines)" >> {log.out}
            echo "Alignment file: $(numfmt --to=iec $ALIGNMENT_SIZE)" >> {log.out}
            
            if [ $CONSENSUS_LINES -gt 1 ]; then
                echo "✓ Consensus contains sequence data" >> {log.out}
            else
                echo "⚠ WARNING: Consensus may be header-only" >> {log.out}
            fi
            echo "===========================" >> {log.out}
        else
            echo "ERROR: Expected output files not found" >> {log.out}
            echo "Expected files:" >> {log.out}
            echo "  Consensus: {output.consensus}" >> {log.out}
            echo "  Alignment: {output.alignment}" >> {log.out}
            exit 1
        fi
        """
        
# ===== CONSENSUS COMBINATION AND LOG RULES =====

# Rename headers in consensus files for merge mode
rule rename_and_combine_cons_merge:
    input:
        consensus_files=get_all_mge_consensus_files("merge")
    output:
        complete=os.path.join(barcode_recovery_dir_merge, "logs/rename_consensus/rename_complete.txt"),
        concat_cons=os.path.join(barcode_recovery_dir_merge, f"consensus/{run_name}_cons_combined-merge.fasta")
    log:
        os.path.join(barcode_recovery_dir_merge, "logs/rename_consensus/rename_fasta.log")
    params:
        script=os.path.join(workflow.basedir, "scripts/rename_headers.py"),
        preprocessing_mode="merge"
    threads: rule_resources["rename_and_combine_cons"]["threads"]
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["rename_and_combine_cons"]["mem_mb"] * attempt
    retries: 3
    shell:
        """ 
        set -euo pipefail

        # Check all input files exist
        for file in {input.consensus_files}; do
            if [ ! -f "$file" ]; then
                echo "Error: Required input file does not exist: $file" >> {log}
                echo "Cannot proceed with renaming and combining consensus files." >> {log}
                echo "If any MitoGeneExtractor jobs failed, please rerun the workflow to generate those files." >> {log}

                exit 1
            fi
        done
        
        # Run script
        python {params.script} \\
            --input-files {input.consensus_files} \\
            --concatenated-consensus {output.concat_cons} \\
            --complete-file={output.complete} \\
            --log-file={log} \\
            --preprocessing-mode={params.preprocessing_mode} \\
            --threads={threads}
    
        # Touch completion file to ensure timestamp is updated
        touch {output.complete}
        """

# Rename headers in consensus files for concat mode
rule rename_and_combine_cons_concat:
    input:
        consensus_files=get_all_mge_consensus_files("concat")
    output:
        complete=os.path.join(barcode_recovery_dir_concat, "logs/rename_consensus/rename_complete.txt"),
        concat_cons=os.path.join(barcode_recovery_dir_concat, f"consensus/{run_name}_cons_combined-concat.fasta")
    log:
        os.path.join(barcode_recovery_dir_concat, "logs/rename_consensus/rename_fasta.log")
    params:
        script=os.path.join(workflow.basedir, "scripts/rename_headers.py"),
        preprocessing_mode="concat"
    threads: rule_resources["rename_and_combine_cons"]["threads"]
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["rename_and_combine_cons"]["mem_mb"] * attempt
    retries: 3
    shell:
        """ 
        set -euo pipefail

        # Check all input files exist
        for file in {input.consensus_files}; do
            if [ ! -f "$file" ]; then
                echo "Error: Required input file does not exist: $file" >> {log}
                echo "Cannot proceed with renaming and combining consensus files." >> {log}
                echo "If any MitoGeneExtractor jobs failed, please rerun the workflow to generate those files." >> {log}
                exit 1
            fi
        done
        
        # Run script
        python {params.script} \\
            --input-files {input.consensus_files} \\
            --concatenated-consensus {output.concat_cons} \\
            --complete-file={output.complete} \\
            --log-file={log} \\
            --preprocessing-mode={params.preprocessing_mode} \\
            --threads={threads}
    
        # Touch completion file to ensure timestamp is updated
        touch {output.complete}
        """

rule gzip_merged_clean:
    input:
        # Wait for all MGE merge jobs across all samples to finish
        merge_complete=os.path.join(barcode_recovery_dir_merge, f"consensus/{run_name}_cons_combined-merge.fasta")
    output:
        completion_flag=os.path.join(barcode_recovery_dir_merge, "logs/gzip_merged_clean_complete.txt")
    log:
        os.path.join(barcode_recovery_dir_merge, "logs/gzip_merged_clean.log")
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["gzip_merged_clean"]["mem_mb"] * attempt
    retries: 2
    run:
        compressed = 0
        skipped = 0
        errors = 0
        
        with open(log[0], 'w') as logf:
            logf.write(f"Starting gzip of merged_clean.fastq files: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            logf.write(f"Searching in: {preprocessing_dir_merge}\n\n")
        
            fastq_files = glob.glob(
                os.path.join(preprocessing_dir_merge, "trimmed_data", "*", "*_merged_clean.fastq")
            )
        
            if not fastq_files:
                logf.write("No *_merged_clean.fastq files found.\n")
            else:
                logf.write(f"Found {len(fastq_files)} file(s) to compress:\n")
                for fastq_file in sorted(fastq_files):
                    # Skip empty placeholder files
                    if os.path.getsize(fastq_file) == 0:
                        logf.write(f"  SKIPPED (empty placeholder): {os.path.basename(fastq_file)}\n")
                        skipped += 1
                        continue
                    try:
                        size_before = os.path.getsize(fastq_file)
                        result = subprocess.run(
                            ["gzip", "-f", fastq_file],
                            capture_output=True, text=True
                        )
                        if result.returncode != 0:
                            raise RuntimeError(result.stderr.strip())
                        gz_file = fastq_file + ".gz"
                        size_after = os.path.getsize(gz_file)
                        saving_pct = 100 * (1 - size_after / size_before) if size_before > 0 else 0
                        logf.write(
                            f"  Compressed: {os.path.basename(fastq_file)} "
                            f"({size_before / 1024**2:.1f} MB -> {size_after / 1024**2:.1f} MB, "
                            f"{saving_pct:.0f}% reduction)\n"
                        )
                        compressed += 1
                    except Exception as e:
                        logf.write(f"  ERROR compressing {fastq_file}: {e}\n")
                        errors += 1
        
            logf.write(f"\nSUMMARY:\n")
            logf.write(f"  Compressed: {compressed}\n")
            logf.write(f"  Skipped (empty): {skipped}\n")
            logf.write(f"  Errors: {errors}\n")
            logf.write(f"  Completed: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        
        if errors > 0:
            raise RuntimeError(f"gzip_merged_clean: {errors} file(s) failed to compress. See {log[0]}")
        
        with open(output.completion_flag, 'w') as f:
            f.write(f"gzip_merged_clean completed at {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write(f"Files compressed: {compressed}\n")
            f.write(f"Files skipped (empty): {skipped}\n")

# Remove exonerate intermediate files after consensus generation
rule remove_exonerate_intermediates:
    input:
        merge_consensus=os.path.join(barcode_recovery_dir_merge, f"consensus/{run_name}_cons_combined-merge.fasta"),
        concat_consensus=os.path.join(barcode_recovery_dir_concat, f"consensus/{run_name}_cons_combined-concat.fasta"),
        merge_rename_complete=os.path.join(barcode_recovery_dir_merge, "logs/rename_consensus/rename_complete.txt"),
        concat_rename_complete=os.path.join(barcode_recovery_dir_concat, "logs/rename_consensus/rename_complete.txt")
    output:
        merge_completion=os.path.join(barcode_recovery_dir_merge, "logs/exonerate_int_cleanup_complete.txt"),
        concat_completion=os.path.join(barcode_recovery_dir_concat, "logs/exonerate_int_cleanup_complete.txt")
    threads: 1
    resources:
        mem_mb=2048
    retries: 2
    run:     
        # Clean each mode directory and write separate completion files
        for mode_dir, mode_name, completion_file in [
            (barcode_recovery_dir_merge, "merge", output.merge_completion),
            (barcode_recovery_dir_concat, "concat", output.concat_completion)
        ]:
            removed_files = 0
            total_size_saved = 0
            
            with open(completion_file, 'w') as log:
                log.write(f"Starting exonerate intermediate file cleanup for {mode_name} mode: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                log.write(f"Mode directory: {mode_dir}\n\n")
                
                exonerate_pattern = os.path.join(mode_dir, "Concatenated_exonerate_input_*")
                exonerate_files = glob.glob(exonerate_pattern)
                
                if not exonerate_files:
                    log.write(f"No Concatenated_exonerate_input_* files found\n")
                else:
                    log.write(f"Found {len(exonerate_files)} files to remove:\n")
                    
                    for file_path in exonerate_files:
                        try:
                            file_size = os.path.getsize(file_path)
                            os.remove(file_path)
                            
                            size_mb = file_size / (1024 * 1024)
                            log.write(f"  - Removed: {os.path.basename(file_path)} ({size_mb:.2f} MB)\n")
                            
                            removed_files += 1
                            total_size_saved += file_size
                            
                        except OSError as e:
                            log.write(f"  - ERROR removing {file_path}: {e}\n")
                
                # Summary
                total_saved_mb = total_size_saved / (1024 * 1024)
                log.write(f"\nCLEANUP SUMMARY for {mode_name} mode:\n")
                log.write(f"  Total files removed: {removed_files}\n")
                log.write(f"  Total disk space saved: {total_saved_mb:.2f} MB\n")
                log.write(f"  Completed: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            
# Create list of alignment files for stats for merge mode
rule create_alignment_log_merge:
    input:
        concat_cons=os.path.join(barcode_recovery_dir_merge, f"consensus/{run_name}_cons_combined-merge.fasta")
    output:
        alignment_log=os.path.join(barcode_recovery_dir_merge, "logs/mge/alignment_files.log")
    params:
        alignment_dir=os.path.join(barcode_recovery_dir_merge, "alignment")
    threads: 1
    resources:
        mem_mb=4096
    retries: 3
    shell:
        "find {params.alignment_dir} -name '*_align_*.fas' -type f 2>/dev/null | sort > {output.alignment_log} || touch {output.alignment_log}"

# Create list of alignment files for stats for concat mode
rule create_alignment_log_concat:
    input:
        concat_cons=os.path.join(barcode_recovery_dir_concat, f"consensus/{run_name}_cons_combined-concat.fasta")
    output:
        alignment_log=os.path.join(barcode_recovery_dir_concat, "logs/mge/alignment_files.log")
    params:
        alignment_dir=os.path.join(barcode_recovery_dir_concat, "alignment")
    threads: 1
    resources:
        mem_mb=4096
    retries: 3
    shell:
        "find {params.alignment_dir} -name '*_align_*.fas' -type f 2>/dev/null | sort > {output.alignment_log} || touch {output.alignment_log}"
                

# ===== FASTA CLEANER RULES - MERGE MODE =====

# Sequential filtering rules (fasta_cleaner)
rule human_cox1_filter_merge:
    input:
        alignment_log=os.path.join(barcode_recovery_dir_merge, "logs/mge/alignment_files.log")
    output:
        filtered_files=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/01_human_filtered/human_filtered.txt"),
        metrics=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/01_human_filtered/human_filter_metrics.csv")
    params:
        script=os.path.join(workflow.basedir, "scripts/01_human_cox1_filter.py"),
        output_dir=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/01_human_filtered"),
        human_threshold=fasta_cleaner_params["human_threshold"]
    log:
        os.path.join(barcode_recovery_dir_merge, "logs/fasta_cleaner/01_human_filter.log")
    threads: rule_resources["human_cox1_filter"]["threads"]
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["human_cox1_filter"]["mem_mb"] * attempt
    retries: 2
    shell:
        """
        set -euo pipefail

        mkdir -p {params.output_dir}
        
        python {params.script} \\
            --input-log {input.alignment_log} \\
            --output-dir {params.output_dir} \\
            --filtered-files-list {output.filtered_files} \\
            --metrics-csv {output.metrics} \\
            --human-threshold {params.human_threshold} \\
            --threads {threads} \\
            2>&1 | tee {log}
        """

rule at_content_filter_merge:
    input:
        human_filtered=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/01_human_filtered/human_filtered.txt")
    output:
        filtered_files=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/02_at_filtered/at_filtered.txt"),
        summary_metrics=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/02_at_filtered/at_filter_summary.csv")
    params:
        script=os.path.join(workflow.basedir, "scripts/02_at_content_filter.py"),
        output_dir=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/02_at_filtered"),
        at_threshold=fasta_cleaner_params["at_difference"],
        at_mode=fasta_cleaner_params["at_mode"],
        consensus_threshold=fasta_cleaner_params["consensus_threshold"]
    log:
        os.path.join(barcode_recovery_dir_merge, "logs/fasta_cleaner/02_at_filter.log")
    threads: rule_resources["at_content_filter"]["threads"]
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["at_content_filter"]["mem_mb"] * attempt
    retries: 2
    shell:
        """
        set -euo pipefail

        mkdir -p {params.output_dir}
        
        python -u {params.script} \\
            --input_files {input.human_filtered} \\
            --output_dir {params.output_dir} \\
            --filtered-files-list {output.filtered_files} \\
            --at_threshold {params.at_threshold} \\
            --at_mode {params.at_mode} \\
            --consensus_threshold {params.consensus_threshold} \\
            --threads {threads} \\
            2>&1 | tee {log}
        """

rule statistical_outlier_filter_merge:
    input:
        filtered_files=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/02_at_filtered/at_filtered.txt")
    output:
        filtered_files=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/03_outlier_filtered/outlier_filtered.txt"),
        summary_metrics=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/03_outlier_filtered/outlier_filter_summary_metrics.csv"),
        individual_metrics=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/03_outlier_filtered/outlier_filter_individual_metrics.csv")
    params:
        script=os.path.join(workflow.basedir, "scripts/03_statistical_outlier_filter.py"),
        output_dir=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/03_outlier_filtered"),
        outlier_percentile=fasta_cleaner_params["outlier_percentile"],
        consensus_threshold=fasta_cleaner_params["consensus_threshold"]
    log:
        os.path.join(barcode_recovery_dir_merge, "logs/fasta_cleaner/03_outlier_filter.log")
    threads: rule_resources["statistical_outlier_filter"]["threads"]
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["statistical_outlier_filter"]["mem_mb"] * attempt
    retries: 2
    shell:
        """
        set -euo pipefail

        mkdir -p {params.output_dir}
        
        python {params.script} \\
            --input-files-list {input.filtered_files} \\
            --output-dir {params.output_dir} \\
            --filtered-files-list {output.filtered_files} \\
            --summary-csv {output.individual_metrics} \\
            --metrics-csv {output.summary_metrics} \\
            --outlier-percentile {params.outlier_percentile} \\
            --consensus-threshold {params.consensus_threshold} \\
            --threads {threads} \\
            2>&1 | tee {log}
        """

if use_reference_filtering:
    # Version WITH reference filtering - MERGE MODE
    rule reference_filter_merge:
        input:
            filtered_files=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/03_outlier_filtered/outlier_filtered.txt")
        output:
            filtered_files=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/04_reference_filtered/reference_filtered.txt"),
            metrics=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/04_reference_filtered/reference_filter_metrics.csv")
        params:
            script=os.path.join(workflow.basedir, "scripts/04_reference_filter.py"),
            output_dir=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/04_reference_filtered"),
            reference_dir=reference_dir,
            filter_mode=fasta_cleaner_params["reference_filter_mode"],
        log:
            os.path.join(barcode_recovery_dir_merge, "logs/fasta_cleaner/04_reference_filter.log")
        threads: rule_resources["reference_filter"]["threads"]
        resources:
            mem_mb=lambda wildcards, attempt: rule_resources["reference_filter"]["mem_mb"] * attempt
        retries: 2
        shell:
            """
            set -euo pipefail

            mkdir -p {params.output_dir}

            python {params.script} \\
                --input-files-list {input.filtered_files} \\
                --output-dir {params.output_dir} \\
                --filtered-files-list {output.filtered_files} \\
                --metrics-csv {output.metrics} \\
                --reference-dir {params.reference_dir} \\
                --filter-mode {params.filter_mode} \\
                --threads {threads} \\
                2>&1 | tee {log}
            """

    rule consensus_generation_merge:
        input:
            filtered_files=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/04_reference_filtered/reference_filtered.txt")
        output:
            concat_consensus=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/cleaned_cons_combined.fasta"),
            consensus_metrics=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/05_cleaned_consensus/cleaned_cons_metrics-merge.csv")
        params:
            script=os.path.join(workflow.basedir, "scripts/05_consensus_generator.py"),
            output_dir=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/05_cleaned_consensus"),
            consensus_threshold=fasta_cleaner_params["consensus_threshold"],
            preprocessing_mode="merge"
        log:
            os.path.join(barcode_recovery_dir_merge, "logs/fasta_cleaner/05_consensus_generation.log")
        threads: rule_resources["consensus_generation"]["threads"]
        resources:
            mem_mb=lambda wildcards, attempt: rule_resources["consensus_generation"]["mem_mb"] * attempt
        retries: 2
        shell:
            """
            set -euo pipefail

            python {params.script} \\
                --input-files-list {input.filtered_files} \\
                --output-dir {params.output_dir} \\
                --consensus-fasta {output.concat_consensus} \\
                --consensus-metrics {output.consensus_metrics} \\
                --consensus-threshold {params.consensus_threshold} \\
                --preprocessing-mode {params.preprocessing_mode} \\
                --threads {threads} \\
                2>&1 | tee {log}
            """

    rule aggregate_filter_metrics_merge:
        input:
            human_metrics=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/01_human_filtered/human_filter_metrics.csv"),
            at_metrics=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/02_at_filtered/at_filter_summary.csv"),
            outlier_metrics=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/03_outlier_filtered/outlier_filter_individual_metrics.csv"),
            reference_metrics=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/04_reference_filtered/reference_filter_metrics.csv"),
            consensus_metrics=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/05_cleaned_consensus/cleaned_cons_metrics-merge.csv")
        output:
            completion_flag=os.path.join(barcode_recovery_dir_merge, "logs/fasta_cleaner_complete.txt"),
            combined_statistics=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/combined_statistics.csv")
        params:
            script=os.path.join(workflow.basedir, "scripts/06_aggregate_filter_metrics.py"),
            output_dir=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner")
        log:
            os.path.join(barcode_recovery_dir_merge, "logs/fasta_cleaner/06_aggregate_metrics.log")
        threads: 2
        resources:
            mem_mb=3076
        retries: 2
        shell:
            """
            set -euo pipefail

            python {params.script} \\
                --human-metrics {input.human_metrics} \\
                --at-metrics {input.at_metrics} \\
                --outlier-metrics {input.outlier_metrics} \\
                --reference-metrics {input.reference_metrics} \\
                --consensus-metrics {input.consensus_metrics} \\
                --output-dir {params.output_dir} \\
                --combined-statistics {output.combined_statistics} \\
                --threads {threads} \\
                2>&1 | tee {log}
                
            touch {output.completion_flag}
            """
else:
    # Version WITHOUT reference filtering
    rule consensus_generation_merge:
        input:
            filtered_files=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/03_outlier_filtered/outlier_filtered.txt")
        output:
            concat_consensus=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/cleaned_cons_combined.fasta"),
            consensus_metrics=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/05_cleaned_consensus/cleaned_cons_metrics-merge.csv")
        params:
            script=os.path.join(workflow.basedir, "scripts/05_consensus_generator.py"),
            output_dir=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/05_cleaned_consensus"),
            consensus_threshold=fasta_cleaner_params["consensus_threshold"],
            preprocessing_mode="merge"
        log:
            os.path.join(barcode_recovery_dir_merge, "logs/fasta_cleaner/05_consensus_generation.log")
        threads: rule_resources["consensus_generation"]["threads"]
        resources:
            mem_mb=lambda wildcards, attempt: rule_resources["consensus_generation"]["mem_mb"] * attempt
        retries: 2
        shell:
            """
            set -euo pipefail

            python {params.script} \\
                --input-files-list {input.filtered_files} \\
                --output-dir {params.output_dir} \\
                --consensus-fasta {output.concat_consensus} \\
                --consensus-metrics {output.consensus_metrics} \\
                --consensus-threshold {params.consensus_threshold} \\
                --preprocessing-mode {params.preprocessing_mode} \\
                --threads {threads} \\
                2>&1 | tee {log}
            """

    rule aggregate_filter_metrics_merge:
        input:
            human_metrics=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/01_human_filtered/human_filter_metrics.csv"),
            at_metrics=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/02_at_filtered/at_filter_summary.csv"),
            outlier_metrics=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/03_outlier_filtered/outlier_filter_individual_metrics.csv"),
            consensus_metrics=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/05_cleaned_consensus/cleaned_cons_metrics-merge.csv")
        output:
            completion_flag=os.path.join(barcode_recovery_dir_merge, "logs/fasta_cleaner_complete.txt"),
            combined_statistics=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/combined_statistics.csv")
        params:
            script=os.path.join(workflow.basedir, "scripts/06_aggregate_filter_metrics.py"),
            output_dir=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner")
        log:
            os.path.join(barcode_recovery_dir_merge, "logs/fasta_cleaner/06_aggregate_metrics.log")
        threads: 2
        resources:
            mem_mb=3076
        retries: 2
        shell:
            """
            set -euo pipefail

            python {params.script} \\
                --human-metrics {input.human_metrics} \\
                --at-metrics {input.at_metrics} \\
                --outlier-metrics {input.outlier_metrics} \\
                --consensus-metrics {input.consensus_metrics} \\
                --output-dir {params.output_dir} \\
                --combined-statistics {output.combined_statistics} \\
                --threads {threads} \\
                2>&1 | tee {log}
                
            touch {output.completion_flag}
            """
            
# ===== FASTA CLEANER RULES - CONCAT MODE =====

# Sequential filtering rules (fasta_cleaner)
rule human_cox1_filter_concat:
    input:
        alignment_log=os.path.join(barcode_recovery_dir_concat, "logs/mge/alignment_files.log")
    output:
        filtered_files=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/01_human_filtered/human_filtered.txt"),
        metrics=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/01_human_filtered/human_filter_metrics.csv")
    params:
        script=os.path.join(workflow.basedir, "scripts/01_human_cox1_filter.py"),
        output_dir=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/01_human_filtered"),
        human_threshold=fasta_cleaner_params["human_threshold"]
    log:
        os.path.join(barcode_recovery_dir_concat, "logs/fasta_cleaner/01_human_filter.log")
    threads: rule_resources["human_cox1_filter"]["threads"]
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["human_cox1_filter"]["mem_mb"] * attempt
    retries: 2
    shell:
        """
        set -euo pipefail

        mkdir -p {params.output_dir}
        
        python {params.script} \\
            --input-log {input.alignment_log} \\
            --output-dir {params.output_dir} \\
            --filtered-files-list {output.filtered_files} \\
            --metrics-csv {output.metrics} \\
            --human-threshold {params.human_threshold} \\
            --threads {threads} \\
            2>&1 | tee {log}
        """

rule at_content_filter_concat:
    input:
        human_filtered=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/01_human_filtered/human_filtered.txt")
    output:
        filtered_files=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/02_at_filtered/at_filtered.txt"),
        summary_metrics=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/02_at_filtered/at_filter_summary.csv")
    params:
        script=os.path.join(workflow.basedir, "scripts/02_at_content_filter.py"),
        output_dir=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/02_at_filtered"),
        at_threshold=fasta_cleaner_params["at_difference"],
        at_mode=fasta_cleaner_params["at_mode"],
        consensus_threshold=fasta_cleaner_params["consensus_threshold"]
    log:
        os.path.join(barcode_recovery_dir_concat, "logs/fasta_cleaner/02_at_filter.log")
    threads: rule_resources["at_content_filter"]["threads"]
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["at_content_filter"]["mem_mb"] * attempt
    retries: 2
    shell:
        """
        set -euo pipefail

        mkdir -p {params.output_dir}
        
        python -u {params.script} \\
            --input_files {input.human_filtered} \\
            --output_dir {params.output_dir} \\
            --filtered-files-list {output.filtered_files} \\
            --at_threshold {params.at_threshold} \\
            --at_mode {params.at_mode} \\
            --consensus_threshold {params.consensus_threshold} \\
            --threads {threads} \\
            2>&1 | tee {log}
        """

rule statistical_outlier_filter_concat:
    input:
        filtered_files=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/02_at_filtered/at_filtered.txt")
    output:
        filtered_files=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/03_outlier_filtered/outlier_filtered.txt"),
        summary_metrics=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/03_outlier_filtered/outlier_filter_summary_metrics.csv"),
        individual_metrics=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/03_outlier_filtered/outlier_filter_individual_metrics.csv")
    params:
        script=os.path.join(workflow.basedir, "scripts/03_statistical_outlier_filter.py"),
        output_dir=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/03_outlier_filtered"),
        outlier_percentile=fasta_cleaner_params["outlier_percentile"],
        consensus_threshold=fasta_cleaner_params["consensus_threshold"]
    log:
        os.path.join(barcode_recovery_dir_concat, "logs/fasta_cleaner/03_outlier_filter.log")
    threads: rule_resources["statistical_outlier_filter"]["threads"]
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["statistical_outlier_filter"]["mem_mb"] * attempt
    retries: 2
    shell:
        """
        set -euo pipefail

        mkdir -p {params.output_dir}
        
        python {params.script} \\
            --input-files-list {input.filtered_files} \\
            --output-dir {params.output_dir} \\
            --filtered-files-list {output.filtered_files} \\
            --summary-csv {output.individual_metrics} \\
            --metrics-csv {output.summary_metrics} \\
            --outlier-percentile {params.outlier_percentile} \\
            --consensus-threshold {params.consensus_threshold} \\
            --threads {threads} \\
            2>&1 | tee {log}
        """
        
if use_reference_filtering:
    # Version WITH reference filtering - CONCAT MODE  
    rule reference_filter_concat:
        input:
            filtered_files=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/03_outlier_filtered/outlier_filtered.txt")
        output:
            filtered_files=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/04_reference_filtered/reference_filtered.txt"),
            metrics=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/04_reference_filtered/reference_filter_metrics.csv")
        params:
            script=os.path.join(workflow.basedir, "scripts/04_reference_filter.py"),
            output_dir=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/04_reference_filtered"),
            reference_dir=reference_dir,
            filter_mode=fasta_cleaner_params["reference_filter_mode"],
        log:
            os.path.join(barcode_recovery_dir_concat, "logs/fasta_cleaner/04_reference_filter.log")
        threads: rule_resources["reference_filter"]["threads"]
        resources:
            mem_mb=lambda wildcards, attempt: rule_resources["reference_filter"]["mem_mb"] * attempt
        retries: 2
        shell:
            """
            set -euo pipefail
            
            mkdir -p {params.output_dir}
            
            python {params.script} \\
                --input-files-list {input.filtered_files} \\
                --output-dir {params.output_dir} \\
                --filtered-files-list {output.filtered_files} \\
                --metrics-csv {output.metrics} \\
                --reference-dir {params.reference_dir} \\
                --filter-mode {params.filter_mode} \\
                --threads {threads} \\
                2>&1 | tee {log}
            """

    rule consensus_generation_concat:
        input:
            filtered_files=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/04_reference_filtered/reference_filtered.txt")
        output:
            concat_consensus=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/cleaned_cons_combined.fasta"),
            consensus_metrics=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/05_cleaned_consensus/cleaned_cons_metrics-concat.csv")
        params:
            script=os.path.join(workflow.basedir, "scripts/05_consensus_generator.py"),
            output_dir=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/05_cleaned_consensus"),
            consensus_threshold=fasta_cleaner_params["consensus_threshold"],
            preprocessing_mode="concat"
        log:
            os.path.join(barcode_recovery_dir_concat, "logs/fasta_cleaner/05_consensus_generation.log")
        threads: rule_resources["consensus_generation"]["threads"]
        resources:
            mem_mb=lambda wildcards, attempt: rule_resources["consensus_generation"]["mem_mb"] * attempt
        retries: 2
        shell:
            """
            set -euo pipefail

            python {params.script} \\
                --input-files-list {input.filtered_files} \\
                --output-dir {params.output_dir} \\
                --consensus-fasta {output.concat_consensus} \\
                --consensus-metrics {output.consensus_metrics} \\
                --consensus-threshold {params.consensus_threshold} \\
                --preprocessing-mode {params.preprocessing_mode} \\
                --threads {threads} \\
                2>&1 | tee {log}
            """

    rule aggregate_filter_metrics_concat:
        input:
            human_metrics=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/01_human_filtered/human_filter_metrics.csv"),
            at_metrics=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/02_at_filtered/at_filter_summary.csv"),
            outlier_metrics=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/03_outlier_filtered/outlier_filter_individual_metrics.csv"),
            reference_metrics=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/04_reference_filtered/reference_filter_metrics.csv"),
            consensus_metrics=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/05_cleaned_consensus/cleaned_cons_metrics-concat.csv")
        output:
            completion_flag=os.path.join(barcode_recovery_dir_concat, "logs/fasta_cleaner_complete.txt"),
            combined_statistics=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/combined_statistics.csv")
        params:
            script=os.path.join(workflow.basedir, "scripts/06_aggregate_filter_metrics.py"),
            output_dir=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner")
        log:
            os.path.join(barcode_recovery_dir_concat, "logs/fasta_cleaner/06_aggregate_metrics.log")
        threads: 2
        resources:
            mem_mb=3076
        retries: 2
        shell:
            """
            set -euo pipefail

            python {params.script} \\
                --human-metrics {input.human_metrics} \\
                --at-metrics {input.at_metrics} \\
                --outlier-metrics {input.outlier_metrics} \\
                --reference-metrics {input.reference_metrics} \\
                --consensus-metrics {input.consensus_metrics} \\
                --output-dir {params.output_dir} \\
                --combined-statistics {output.combined_statistics} \\
                --threads {threads} \\
                2>&1 | tee {log}
                
            touch {output.completion_flag}
            """

else:
    # Version WITHOUT reference filtering
    rule consensus_generation_concat:
        input:
            filtered_files=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/03_outlier_filtered/outlier_filtered.txt")
        output:
            concat_consensus=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/cleaned_cons_combined.fasta"),
            consensus_metrics=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/05_cleaned_consensus/cleaned_cons_metrics-concat.csv")
        params:
            script=os.path.join(workflow.basedir, "scripts/05_consensus_generator.py"),
            output_dir=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/05_cleaned_consensus"),
            consensus_threshold=fasta_cleaner_params["consensus_threshold"],
            preprocessing_mode="concat"
        log:
            os.path.join(barcode_recovery_dir_concat, "logs/fasta_cleaner/05_consensus_generation.log")
        threads: rule_resources["consensus_generation"]["threads"]
        resources:
            mem_mb=lambda wildcards, attempt: rule_resources["consensus_generation"]["mem_mb"] * attempt
        retries: 2
        shell:
            """
            set -euo pipefail

            python {params.script} \\
                --input-files-list {input.filtered_files} \\
                --output-dir {params.output_dir} \\
                --consensus-fasta {output.concat_consensus} \\
                --consensus-metrics {output.consensus_metrics} \\
                --consensus-threshold {params.consensus_threshold} \\
                --preprocessing-mode {params.preprocessing_mode} \\
                --threads {threads} \\
                2>&1 | tee {log}
            """

    rule aggregate_filter_metrics_concat:
        input:
            human_metrics=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/01_human_filtered/human_filter_metrics.csv"),
            at_metrics=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/02_at_filtered/at_filter_summary.csv"),
            outlier_metrics=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/03_outlier_filtered/outlier_filter_individual_metrics.csv"),
            consensus_metrics=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/05_cleaned_consensus/cleaned_cons_metrics-concat.csv")
        output:
            completion_flag=os.path.join(barcode_recovery_dir_concat, "logs/fasta_cleaner_complete.txt"),
            combined_statistics=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/combined_statistics.csv")
        params:
            script=os.path.join(workflow.basedir, "scripts/06_aggregate_filter_metrics.py"),
            output_dir=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner")
        log:
            os.path.join(barcode_recovery_dir_concat, "logs/fasta_cleaner/06_aggregate_metrics.log")
        threads: 1
        resources:
            mem_mb=3076
        retries: 2
        shell:
            """
            set -euo pipefail

            python {params.script} \\
                --human-metrics {input.human_metrics} \\
                --at-metrics {input.at_metrics} \\
                --outlier-metrics {input.outlier_metrics} \\
                --consensus-metrics {input.consensus_metrics} \\
                --output-dir {params.output_dir} \\
                --combined-statistics {output.combined_statistics} \\
                --threads {threads} \\
                2>&1 | tee {log}
                
            touch {output.completion_flag}
            """
            
# ===== CLEANUP AND CONSENSUS COMBINATION RULES =====

# Clean up intermediate fasta files from filtering directories (remove intermediate FASTA files)
rule remove_fasta_cleaner_files:
    input:
        # Wait for filtering to complete
        merge_cleaning_complete=os.path.join(barcode_recovery_dir_merge, "logs/fasta_cleaner_complete.txt"),
        concat_cleaning_complete=os.path.join(barcode_recovery_dir_concat, "logs/fasta_cleaner_complete.txt"),
        # Ensure final consensus sequences are created before cleanup
        merge_consensus=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/cleaned_cons_combined.fasta"),
        concat_consensus=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/cleaned_cons_combined.fasta")
    output:
        merge_cleanup_complete=os.path.join(barcode_recovery_dir_merge, "logs/fasta_cleaner/fasta_cleaner_complete.txt"),
        concat_cleanup_complete=os.path.join(barcode_recovery_dir_concat, "logs/fasta_cleaner/fasta_cleaner_complete.txt")
    params:
        merge_dir=barcode_recovery_dir_merge,
        concat_dir=barcode_recovery_dir_concat,
        use_ref_filtering=str(use_reference_filtering)
    threads: 1
    resources:
        mem_mb=2048
    retries: 2
    run:
        def cleanup_fasta_files(mode_dir, mode_name, completion_file):       
            removed_files = 0
            total_size_saved = 0
            
            with open(completion_file, 'w') as log:
                log.write(f"Starting intermediate fasta cleanup for {mode_name} mode: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                log.write(f"Mode directory: {mode_dir}\n")
                log.write(f"Reference filtering enabled: {params.use_ref_filtering}\n\n")
                
                # Define directories to clean
                filter_dirs = [
                    "01_human_filtered",
                    "02_at_filtered", 
                    "03_outlier_filtered"
                ]
                
                # Add reference filtering directory if enabled
                if params.use_ref_filtering == "True":
                    filter_dirs.append("04_reference_filtered")
                    log.write("Including 04_reference_filtered directory in cleanup\n")
                else:
                    log.write("Skipping 04_reference_filtered directory (not enabled)\n")
                
                log.write(f"Directories to clean: {', '.join(filter_dirs)}\n\n")
                
                # Clean each filtering directory
                for filter_dir in filter_dirs:
                    dir_path = os.path.join(mode_dir, "fasta_cleaner", filter_dir)
                    
                    if not os.path.exists(dir_path):
                        log.write(f"[{filter_dir}] Directory not found: {dir_path}\n")
                        continue
                    
                    log.write(f"[{filter_dir}] Cleaning directory: {dir_path}\n")
                    
                    # Find all .fasta and .fas files
                    fasta_patterns = [
                        os.path.join(dir_path, "*.fasta"),
                        os.path.join(dir_path, "*.fas"),
                        os.path.join(dir_path, "**", "*.fasta"),
                        os.path.join(dir_path, "**", "*.fas")
                    ]
                    
                    fasta_files = []
                    for pattern in fasta_patterns:
                        fasta_files.extend(glob.glob(pattern, recursive=True))
                    
                    # Remove duplicates
                    fasta_files = list(set(fasta_files))
                    
                    if not fasta_files:
                        log.write(f"[{filter_dir}] No .fasta or .fas files found\n")
                        continue
                    
                    log.write(f"[{filter_dir}] Found {len(fasta_files)} fasta files to remove:\n")
                    
                    dir_removed = 0
                    dir_size_saved = 0
                    
                    for fasta_file in fasta_files:
                        try:
                            # Get file size before removal
                            file_size = os.path.getsize(fasta_file)
                            
                            # Remove the file
                            os.remove(fasta_file)
                            
                            # Log the removal
                            size_mb = file_size / (1024 * 1024)
                            log.write(f"  - Removed: {os.path.basename(fasta_file)} ({size_mb:.2f} MB)\n")
                            
                            dir_removed += 1
                            dir_size_saved += file_size
                            
                        except OSError as e:
                            log.write(f"  - ERROR removing {fasta_file}: {e}\n")
                    
                    removed_files += dir_removed
                    total_size_saved += dir_size_saved
                    
                    size_saved_mb = dir_size_saved / (1024 * 1024)
                    log.write(f"[{filter_dir}] Removed {dir_removed} files, saved {size_saved_mb:.2f} MB\n\n")
                
                # Summary
                total_saved_mb = total_size_saved / (1024 * 1024)
                log.write(f"CLEANUP SUMMARY for {mode_name} mode:\n")
                log.write(f"  Total files removed: {removed_files}\n")
                log.write(f"  Total disk space saved: {total_saved_mb:.2f} MB\n")
                log.write(f"  Completed: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                
                return removed_files, total_saved_mb
        
        # Clean up merge mode - write directly to completion file
        merge_removed, merge_saved = cleanup_fasta_files(params.merge_dir, "merge", output.merge_cleanup_complete)
        
        # Clean up concat mode - write directly to completion file
        concat_removed, concat_saved = cleanup_fasta_files(params.concat_dir, "concat", output.concat_cleanup_complete)


def cat_consensus_inputs(wildcards):
    """Consensus FASTAs to concatenate into the run-wide combined file.
    PE: concat + merge pre-clean and post-clean. SE: se pre-clean and post-clean.
    """
    if run_mode == "SE":
        return [
            os.path.join(barcode_recovery_dir_se, f"consensus/{run_name}_cons_combined-se.fasta"),
            os.path.join(barcode_recovery_dir_se, "fasta_cleaner/cleaned_cons_combined.fasta"),
        ]
    return [
        os.path.join(barcode_recovery_dir_concat, f"consensus/{run_name}_cons_combined-concat.fasta"),
        os.path.join(barcode_recovery_dir_merge, f"consensus/{run_name}_cons_combined-merge.fasta"),
        os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/cleaned_cons_combined.fasta"),
        os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/cleaned_cons_combined.fasta"),
    ]

rule cat_consensus_files:
    input:
        cat_consensus_inputs
    output:
        final_consensus=os.path.join(barcode_recovery_dir, f"{run_name}_all_cons_combined.fasta")
    log:
        os.path.join(barcode_recovery_dir, "logs/cat_consensus_files.log")
    threads: 2
    resources:
        mem_mb=3076
    retries: 2
    shell:
        """
        set -euo pipefail
                        
        echo "Concatenating consensus multi-FASTA files for {run_name}" > {log}
        echo "Started: $(date)" >> {log}

        # Concatenation (2 files for SE, 4 for PE)
        cat {input} > {output.final_consensus}

        echo "Completed: $(date)" >> {log}
        echo "Output file size: $(numfmt --to=iec $(stat -c%s {output.final_consensus}))" >> {log}
        """
        
rule barcode_consensus_count:
    input:
        all_cons_combined=os.path.join(barcode_recovery_dir, f"{run_name}_all_cons_combined.fasta")
    output:
        tsv_barcode_count=os.path.join(barcode_recovery_dir, "barcode_consensus_count.tsv"),
        log_barcode_count=os.path.join(barcode_recovery_dir, "logs/barcode_consensus_count.log")
    params:
        script=os.path.join(workflow.basedir, "scripts/barcode_consensus_count.py"),
        outlier_percentile=fasta_cleaner_params["outlier_percentile"],
        consensus_threshold=fasta_cleaner_params["consensus_threshold"]
    threads: 1
    resources:
        mem_mb=1024
    retries: 3
    shell:
        """
        set -euo pipefail
        python {params.script} \\
            --input {input.all_cons_combined} \\
            --tsv {output.tsv_barcode_count} \\
            --log {output.log_barcode_count}
        """
        
        
# ===== VALIDATION RULES =====

def structural_validation_fastas(wildcards):
    """Consensus FASTAs fed to structural validation.
    PE: concat+merge pre/post-clean (4 files). SE: se pre/post-clean (2 files).
    """
    if run_mode == "SE":
        return [
            os.path.join(barcode_recovery_dir_se, f"consensus/{run_name}_cons_combined-se.fasta"),
            os.path.join(barcode_recovery_dir_se, "fasta_cleaner/cleaned_cons_combined.fasta"),
        ]
    return [
        os.path.join(barcode_recovery_dir_concat, f"consensus/{run_name}_cons_combined-concat.fasta"),
        os.path.join(barcode_recovery_dir_merge, f"consensus/{run_name}_cons_combined-merge.fasta"),
        os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/cleaned_cons_combined.fasta"),
        os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/cleaned_cons_combined.fasta"),
    ]

def structural_validation_stats(wildcards):
    """combined_statistics.csv dependency/dependencies, mode-aware (ensures
    the fasta_cleaner chain has completed before validation runs)."""
    if run_mode == "SE":
        return [os.path.join(barcode_recovery_dir_se, "fasta_cleaner/combined_statistics.csv")]
    return [
        os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/combined_statistics.csv"),
        os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/combined_statistics.csv"),
    ]

rule structural_validation:
    input:
        fastas=structural_validation_fastas,
        stats=structural_validation_stats,
        hmm_file=hmm_file_path
    output:
        csv=os.path.join(main_output_dir, f"04_barcode_validation/structural/structural_validation.csv"),
        all_passing_barcode=os.path.join(main_output_dir, f"04_barcode_validation/structural/output_barcode_all_passing.fasta")
    params:
        script=os.path.join(workflow.basedir, "scripts/structural_validation.py"),
        target_marker=structural_validation_params["target"],
        genetic_code=C,
        verbosity=" ".join(["--verbose" if structural_validation_params["verbose"] else ""]).strip()
    log:
        os.path.join(main_output_dir, f"04_barcode_validation/logs/structural_validation.log")
    threads: rule_resources["structural_validation"]["threads"]
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["structural_validation"]["mem_mb"] * attempt,
        slurm_partition=rule_resources["structural_validation"]["partition"]
    retries: 3
    shell:
        """
        set -euo pipefail

        # Ensure HMM file exists
        if [ ! -f "{input.hmm_file}" ]; then
            echo "ERROR: HMM file not found: {input.hmm_file}" >> {log}
            echo "Please ensure the HMM file for target '{params.target_marker}' exists at the specified path." >> {log}
            exit 1
        fi

        python {params.script} \
            --output-csv {output.csv} \
            --input {input.fastas} \
            --hmm {input.hmm_file} \
            --code {params.genetic_code} \
            --threads {threads} \
            --log-file {log} \
            --disable-selection \
            {params.verbosity}
        """

rule local_blast:
    input:
        barcode_fasta=os.path.join(main_output_dir, f"04_barcode_validation/structural/output_barcode_all_passing.fasta")
    output:
        csv=os.path.join(main_output_dir, f"04_barcode_validation/taxonomic/metrics/01_local_blast_output.csv")
    params:
        script=os.path.join(workflow.basedir, "scripts/tv_local_blast.py"),
        database=taxonomic_validation_params["blast_db"],
        verbosity="--verbose" if taxonomic_validation_params["verbose"] else "",
        out_dir=os.path.join(main_output_dir, f"04_barcode_validation/taxonomic/")
    log:
        os.path.join(main_output_dir, f"04_barcode_validation/logs/01_local_blast.log")
    threads: rule_resources["taxonomic_validation"]["threads"]
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["taxonomic_validation"]["mem_mb"] * attempt,
        slurm_partition=rule_resources["taxonomic_validation"]["partition"]
    retries: 3
    shell:
        """
        set -euo pipefail

        # Create output directory
        mkdir -p {params.out_dir}

        # Check if database exists
        if [ ! -e "{params.database}" ]; then
            echo "ERROR: Database not found: {params.database}" >> {log}
            echo "Please ensure the database path exists and contains BLAST database files or a FASTA file." >> {log}
            exit 1
        fi

        echo "Starting taxonomic validation: $(date)" > {log}
        echo "Input FASTA: {input.barcode_fasta}" >> {log}
        echo "Database: {params.database}" >> {log}
        echo "Output directory: {params.out_dir}" >> {log}
        echo "Threads: {threads}" >> {log}
        echo "Verbose: {params.verbosity}" >> {log}

        # If no barcodes passed structural validation the FASTA is empty.
        # BLAST cannot run on an empty query — write a header-only TSV and skip.
        if [ ! -s "{input.barcode_fasta}" ]; then
            echo "No passing barcodes from structural validation — skipping BLAST, writing empty output" >> {log}
            printf 'qseqid\tsseqid\tpident\tlength\tmismatch\tgapopen\tqstart\tqend\tsstart\tsend\tevalue\tbitscore\tstitle\n' > {output.csv}
        else
            python {params.script} \\
                --input {input.barcode_fasta} \\
                --database {params.database} \\
                --output {params.out_dir} \\
                --output-csv {output.csv} \\
                --processes {threads} \\
                {params.verbosity} \\
                >> {log} 2>&1
        fi

        echo "Taxonomic validation completed: $(date)" >> {log}
        """

rule blast2taxonomy:
    input:
        blast_csv=os.path.join(main_output_dir, f"04_barcode_validation/taxonomic/metrics/01_local_blast_output.csv"),
        blast_taxonomy=taxonomic_validation_params["db_taxonomy"],
        expected_taxonomy=taxonomic_validation_params["expected_taxonomy"],
        input_fasta=os.path.join(main_output_dir, f"04_barcode_validation/structural/output_barcode_all_passing.fasta")
    output:
        taxonomy_csv=os.path.join(main_output_dir, f"04_barcode_validation/taxonomic/metrics/02_taxonomic_validation.csv"),
        barcode_fasta=os.path.join(main_output_dir, f"04_barcode_validation/taxonomic/validated_barcodes.fasta"),
        final_fasta=os.path.join(main_output_dir, f"{run_name}_validated_barcodes.fasta")
    params:
        script=os.path.join(workflow.basedir, "scripts/tv_blast2taxonomy.py"),
        taxval_rank=taxonomic_validation_params.get("taxval_rank", "family"),
        min_pident=taxonomic_validation_params.get("min_pident", "80"),
        min_length=taxonomic_validation_params.get("min_length", "100")
    log:
        os.path.join(main_output_dir, f"04_barcode_validation/logs/02_taxonomic_validation.log")
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["blast2taxonomy"]["mem_mb"] * attempt
    retries: 3
    shell:
        """
        set -euo pipefail

        echo "Starting taxonomy validation: $(date)" > {log}
        echo "BLAST CSV: {input.blast_csv}" >> {log}
        echo "BLAST taxonomy: {input.blast_taxonomy}" >> {log}
        echo "Expected taxonomy: {input.expected_taxonomy}" >> {log}
        echo "Input FASTA: {input.input_fasta}" >> {log}
        echo "Validation rank: {params.taxval_rank}" >> {log}
        echo "Output CSV: {output.taxonomy_csv}" >> {log}
        echo "Output FASTA: {output.barcode_fasta}" >> {log}

        # If BLAST produced no hits (header-only CSV), skip taxonomy validation
        # and write empty outputs so the pipeline completes with all samples as FAIL.
        blast_lines=$(wc -l < "{input.blast_csv}")
        if [ "$blast_lines" -le 1 ]; then
            echo "No BLAST results — skipping taxonomy validation, writing empty outputs" >> {log}
            printf 'seq_id,Process_ID,taxval_rank,expected_taxonomy,expected_taxonomy_rank,obs_taxonomy,match_taxonomy,matched_rank,top_matching_hit,pident,length,mismatch,gaps,evalue,selected\n' > {output.taxonomy_csv}
            touch {output.barcode_fasta}
            touch {output.final_fasta}
        else
            python {params.script} \\
                --taxval-rank {params.taxval_rank} \\
                --input-blast-csv {input.blast_csv} \\
                --input-taxonomy-file {input.blast_taxonomy} \\
                --input-exp-taxonomy {input.expected_taxonomy} \\
                --input-fasta {input.input_fasta} \\
                --output-csv {output.taxonomy_csv} \\
                --output-fasta {output.barcode_fasta} \\
                --min-pident {params.min_pident} \\
                --min-length {params.min_length} \\
                --log {log} \\
                >> {log} 2>&1

            # Copy and rename the barcode_fasta to another directory
            cp {output.barcode_fasta} {output.final_fasta}
        fi

        echo "Taxonomic validation completed: $(date)" >> {log}
        """
            
# ===== STATISTICS AND FINAL RULES =====

# Extract, aggregate and calculate stats from inputs for merge mode
rule extract_stats_to_csv_merge:
    input:
        alignment_log=os.path.join(barcode_recovery_dir_merge, "logs/mge/alignment_files.log"),
        out_files=expand(
            os.path.join(barcode_recovery_dir_merge, "out/{sample}_r_{r}_s_{s}_summary.out"),
            sample=list(samples.keys()),
            r=r,
            s=s
        ),
        pre_fasta=os.path.join(barcode_recovery_dir_merge, f"consensus/{run_name}_cons_combined-merge.fasta"),
        cleaning_csv=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/combined_statistics.csv"),
        cleaner_complete=os.path.join(barcode_recovery_dir_merge, "logs/fasta_cleaner_complete.txt"),
        post_fasta=os.path.join(barcode_recovery_dir_merge, "fasta_cleaner/cleaned_cons_combined.fasta"),
        ref_seqs=os.path.join(main_output_dir, "02_references/sequence_references.csv") if run_gene_fetch else []
    output:
        stats=os.path.join(barcode_recovery_dir_merge, f"{run_name}_merge-stats.csv")
    log:
        os.path.join(barcode_recovery_dir_merge, "logs/compile_barcoding_stats.log")
    params:
        script=os.path.join(workflow.basedir, "scripts/compile_barcoding_stats.py"),
        out_dir=os.path.join(barcode_recovery_dir_merge, "out/"),
        ref_seqs_arg=lambda wildcards, input: f"--ref_seqs {input.ref_seqs}" if run_gene_fetch else ""
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["extract_stats_to_csv"]["mem_mb"] * attempt
    retries: 3
    shell:
        """
        set -euo pipefail
        # Run script
        python {params.script} -a {input.alignment_log} -o {output.stats} -od {params.out_dir} -c {input.cleaning_csv} {params.ref_seqs_arg} --pre-fasta {input.pre_fasta} --post-fasta {input.post_fasta} --log-file {log}
        """

rule extract_stats_to_csv_concat:
    input:
        alignment_log=os.path.join(barcode_recovery_dir_concat, "logs/mge/alignment_files.log"),
        out_files=expand(
            os.path.join(barcode_recovery_dir_concat, "out/{sample}_r_{r}_s_{s}_summary.out"),
            sample=list(samples.keys()),
            r=r,
            s=s
        ),
        pre_fasta=os.path.join(barcode_recovery_dir_concat, f"consensus/{run_name}_cons_combined-concat.fasta"),
        cleaning_csv=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/combined_statistics.csv"),
        cleaner_complete=os.path.join(barcode_recovery_dir_concat, "logs/fasta_cleaner_complete.txt"),
        post_fasta=os.path.join(barcode_recovery_dir_concat, "fasta_cleaner/cleaned_cons_combined.fasta"),
        ref_seqs=os.path.join(main_output_dir, "02_references/sequence_references.csv") if run_gene_fetch else []
    output:
        stats=os.path.join(barcode_recovery_dir_concat, f"{run_name}_concat-stats.csv")
    log:
        os.path.join(barcode_recovery_dir_concat, "logs/compile_barcoding_stats.log")
    params:
        script=os.path.join(workflow.basedir, "scripts/compile_barcoding_stats.py"),
        out_dir=os.path.join(barcode_recovery_dir_concat, "out/"),
        ref_seqs_arg=lambda wildcards, input: f"--ref_seqs {input.ref_seqs}" if run_gene_fetch else ""
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["extract_stats_to_csv"]["mem_mb"] * attempt
    retries: 3
    shell:
        """
        set -euo pipefail
        # Run script
        python {params.script} -a {input.alignment_log} -o {output.stats} -od {params.out_dir} -c {input.cleaning_csv} {params.ref_seqs_arg} --pre-fasta {input.pre_fasta} --post-fasta {input.post_fasta} --log-file {log}
        """

def combine_stats_inputs(wildcards):
    """Per-mode stats CSVs to combine. PE: merge + concat. SE: se only."""
    if run_mode == "SE":
        return [os.path.join(barcode_recovery_dir_se, f"{run_name}_se-stats.csv")]
    return [
        os.path.join(barcode_recovery_dir_merge, f"{run_name}_merge-stats.csv"),
        os.path.join(barcode_recovery_dir_concat, f"{run_name}_concat-stats.csv"),
    ]

# Concatenate per-mode stats CSV files with column alignment
rule combine_stats_files:
    input:
        combine_stats_inputs
    output:
        combined_stats=os.path.join(barcode_recovery_dir, f"{run_name}_barcode_recovery_metrics.csv")
    params:
        script=os.path.join(workflow.basedir, "scripts/csv_combiner_mge.py")
    threads: 2
    resources:
        mem_mb=3076
    retries: 3
    shell:
        """
        set -euo pipefail

        python {params.script} \
            --input {input} \
            --output {output.combined_stats}
        """    

rule val_barcode_recovery_metrics_merge:
    input:
        structval=os.path.join(main_output_dir, f"04_barcode_validation/structural/structural_validation.csv"),
        taxval=os.path.join(main_output_dir, f"04_barcode_validation/taxonomic/metrics/02_taxonomic_validation.csv"),
        beegees_csv=os.path.join(barcode_recovery_dir, f"{run_name}_barcode_recovery_metrics.csv")
    output:
        final_metrics=os.path.join(main_output_dir, f"{run_name}_final_metrics.csv")
    params:
        script=os.path.join(workflow.basedir, "scripts/val_csv_merger.py")
    threads: 1
    resources:
        mem_mb=3076
    retries: 3
    shell:
        """
        set -euo pipefail

        python {params.script} \
            --structval-csv {input.structval} \
            --taxval-csv {input.taxval} \
            --beegees-csv {input.beegees_csv} \
            --output {output.final_metrics}
        """    

rule barcoding_outcome:
    input:
        final_metrics=os.path.join(main_output_dir, f"{run_name}_final_metrics.csv"),
        expected_taxonomy=taxonomic_validation_params["expected_taxonomy"]
    output:
        outcome_tsv=os.path.join(main_output_dir, "05_barcoding_outcome/barcoding_outcome.tsv")
    params:
        script=os.path.join(workflow.basedir, "scripts/barcoding_outcome.py")
    log:
        os.path.join(main_output_dir, "05_barcoding_outcome/barcoding_outcome.log")
    threads: 1
    resources:
        mem_mb=4096
    retries: 2
    shell:
        """
        set -euo pipefail

        mkdir -p $(dirname {output.outcome_tsv})

        python {params.script} \
            --metric {input.final_metrics} \
            --taxonomy {input.expected_taxonomy} \
            --output {output.outcome_tsv} \
            --log {log}
        """
    
    
# ===== FASTP SUMMARY RULES =====

rule fastp_summary_concat:
    input:
        json_files=expand(
            os.path.join(preprocessing_dir_concat, "trimmed_data/{sample}/{sample}_fastp_report.json"),
            sample=list(samples.keys())
        )
    output:
        summary_csv=os.path.join(preprocessing_dir_concat, "fastp_summary-concat.csv")
    params:
        script=os.path.join(workflow.basedir, "scripts/fastp_summary_parser.py"),
        input_dir=os.path.join(preprocessing_dir_concat, "trimmed_data")
    log:
        os.path.join(preprocessing_dir_concat, "logs/fastp_summary_concat.log")
    threads: 1
    resources:
        mem_mb=4096
    retries: 2
    shell:
        """
        set -euo pipefail

        mkdir -p $(dirname {log})

        echo "Generating fastp summary CSV for concat mode: $(date)" > {log}
        echo "Input directory: {params.input_dir}" >> {log}
        echo "Output CSV: {output.summary_csv}" >> {log}

        python {params.script} \
            -i {params.input_dir} \
            -o {output.summary_csv} \
            >> {log} 2>&1

        echo "Completed: $(date)" >> {log}
        """

rule fastp_summary_merge:
    input:
        json_files=expand(
            os.path.join(preprocessing_dir_merge, "trimmed_data/{sample}/{sample}_fastp_report.json"),
            sample=list(samples.keys())
        )
    output:
        summary_csv=os.path.join(preprocessing_dir_merge, "fastp_summary-merge.csv")
    params:
        script=os.path.join(workflow.basedir, "scripts/fastp_summary_parser.py"),
        input_dir=os.path.join(preprocessing_dir_merge, "trimmed_data")
    log:
        os.path.join(preprocessing_dir_merge, "logs/fastp_summary_merge.log")
    threads: 1
    resources:
        mem_mb=4096
    retries: 2
    shell:
        """
        set -euo pipefail
    
        mkdir -p $(dirname {log})
    
        echo "Generating fastp summary CSV for merge mode: $(date)" > {log}
        echo "Input directory: {params.input_dir}" >> {log}
        echo "Output CSV: {output.summary_csv}" >> {log}
    
        python {params.script} \
            -i {params.input_dir} \
            -o {output.summary_csv} \
            >> {log} 2>&1
    
        echo "Completed: $(date)" >> {log}
        """
    
# Download NCBI taxdump if needed
rule download_taxdump:
    output:
        nodes = os.path.join(taxdump_dir, "nodes.dmp"),
        names = os.path.join(taxdump_dir, "names.dmp")
    params:
        url        = "https://ftp.ncbi.nlm.nih.gov/pub/taxonomy/taxdump.tar.gz",
        taxdump_dir = taxdump_dir
    log:
        os.path.join(taxdump_dir, "download_taxdump.log")
    resources:
        mem_mb = rule_resources["download_taxdump"]["mem_mb"]
    retries: 2
    shell:
        """
        set -euo pipefail
    
        mkdir -p {params.taxdump_dir}
    
        # Check if nodes.dmp and names.dmp already exist - skip download if so
        if [ -f "{output.nodes}" ] && [ -f "{output.names}" ]; then
            echo "nodes.dmp and names.dmp already exist - skipping download" > {log}
            exit 0
        fi
    
        echo "Downloading taxdump.tar.gz: $(date)" > {log}
        echo "URL: {params.url}" >> {log}
    
        wget -q --show-progress \
            -O {params.taxdump_dir}/taxdump.tar.gz \
            {params.url} \
            >> {log} 2>&1
    
        echo "Download complete: $(date)" >> {log}
        echo "Decompressing..." >> {log}
    
        tar -xzf {params.taxdump_dir}/taxdump.tar.gz \
            -C {params.taxdump_dir} \
            >> {log} 2>&1
    
        echo "Decompression complete: $(date)" >> {log}
    
        # Verify required files exist
        if [ ! -f "{output.nodes}" ]; then
            echo "ERROR: nodes.dmp not found after extraction" >> {log}
            exit 1
        fi
        if [ ! -f "{output.names}" ]; then
            echo "ERROR: names.dmp not found after extraction" >> {log}
            exit 1
        fi
    
        # Remove the tarball to save space
        rm {params.taxdump_dir}/taxdump.tar.gz
        echo "Cleaned up tarball: $(date)" >> {log}
        echo "Taxdump download and extraction complete." >> {log}
        """
    
# MultiQC data & static PNG plot creation
rule multiqc_plots:
    input:
        input_csv        = os.path.join(main_output_dir, "02_references/sequence_references.csv")
                           if run_gene_fetch else config["sequence_reference_file"],
        input_outcome    = os.path.join(main_output_dir, "05_barcoding_outcome/barcoding_outcome.tsv"),
        fastp_concat_csv = (os.path.join(preprocessing_dir_concat, "fastp_summary-concat.csv")
                            if run_mode == "PE"
                            else os.path.join(preprocessing_dir_se, "fastp_summary-se.csv")),
        fastp_merge_csv  = (os.path.join(preprocessing_dir_merge, "fastp_summary-merge.csv")
                            if run_mode == "PE"
                            else os.path.join(preprocessing_dir_se, "fastp_summary-se.csv")),
        final_metrics    = os.path.join(main_output_dir, f"{run_name}_final_metrics.csv"),
        nodes_dmp        = os.path.join(taxdump_dir, "nodes.dmp"),
        names_dmp        = os.path.join(taxdump_dir, "names.dmp")
    output:
        plots_dir_flag = touch(os.path.join(main_output_dir, "05_barcoding_outcome/plots/plots_complete.txt")),
        mqc_dir_flag = touch(os.path.join(main_output_dir, "05_barcoding_outcome/multiqc_report/mqc_in_data/mqc_complete.txt"))
    params:
        script      = os.path.join(workflow.basedir, "scripts/multiqc_plots.R"),
        taxdump_dir = taxdump_dir,
        plots_dir   = os.path.join(main_output_dir, "05_barcoding_outcome/plots"),
        mqc_dir     = os.path.join(main_output_dir, "05_barcoding_outcome/multiqc_report/mqc_in_data"),
        run_mode    = run_mode
    log:
        os.path.join(main_output_dir, "05_barcoding_outcome/logs/multiqc_plots.log")
    resources:
        mem_mb = rule_resources["multiqc_plots"]["mem_mb"]
    retries: 2
    shell:
        """
        set -euo pipefail

        mkdir -p {params.plots_dir}
        mkdir -p {params.mqc_dir}
        mkdir -p $(dirname {log})

        echo "Running multiqc_plots.R: $(date)" > {log}
        echo "Input CSV:        {input.input_csv}" >> {log}
        echo "Input outcome:    {input.input_outcome}" >> {log}
        echo "Fastp concat CSV: {input.fastp_concat_csv}" >> {log}
        echo "Fastp merge CSV:  {input.fastp_merge_csv}" >> {log}
        echo "Final metrics:    {input.final_metrics}" >> {log}
        echo "Taxdump dir:      {params.taxdump_dir}" >> {log}
        echo "Plots output:     {params.plots_dir}" >> {log}
        echo "MQC output:       {params.mqc_dir}" >> {log}

        Rscript {params.script} \
            {input.input_csv} \
            {input.input_outcome} \
            {params.mqc_dir} \
            {input.fastp_concat_csv} \
            {input.fastp_merge_csv} \
            {input.final_metrics} \
            {params.taxdump_dir} \
            {params.run_mode} \
            >> {log} 2>&1
        
        # Move static PNGs from mqc_dir to plots_dir
        find {params.mqc_dir} -maxdepth 1 -name "*.png" \
            -exec mv {{}} {params.plots_dir}/ \;
        
        echo "Completed: $(date)" >> {log}
        """


# Generate MultiQC report
rule multiqc:
    input:
        mqc_flag   = os.path.join(main_output_dir, "05_barcoding_outcome/multiqc_report/mqc_in_data/mqc_complete.txt"),
        mqc_config = os.path.join(workflow.basedir, "../config/multiqc_config.yaml")
    output:
        report = os.path.join(main_output_dir, "05_barcoding_outcome/multiqc_report/multiqc_report.html")
    params:
        mqc_data_dir = os.path.join(main_output_dir, "05_barcoding_outcome/multiqc_report/mqc_in_data"),
        output_dir = os.path.join(main_output_dir, "05_barcoding_outcome/multiqc_report")
    log:
        os.path.join(main_output_dir, "05_barcoding_outcome/logs/multiqc.log")
    resources:
        mem_mb = rule_resources["multiqc"]["mem_mb"]
    retries: 2
    shell:
        """
        set -euo pipefail
        
        mkdir -p {params.output_dir}
        mkdir -p $(dirname {log})
        
        echo "Running MultiQC: $(date)" > {log}
        echo "Data dir:    {params.mqc_data_dir}" >> {log}
        echo "Config:      {input.mqc_config}" >> {log}
        echo "Output dir:  {params.output_dir}" >> {log}
        
        multiqc {params.mqc_data_dir} \
            --config {input.mqc_config} \
            --outdir {params.output_dir} \
            --filename multiqc_report.html \
            --force \
            >> {log} 2>&1
        
        # Copy report back to main output directory
        cp {params.output_dir}/multiqc_report.html \
        {main_output_dir}/multiqc_report.html
        
        echo "MultiQC complete: $(date)" >> {log}
        """

# Final clean up superfluous files
rule cleanup_files:
    input:
        # Merge mode inputs
        merge_summary_csv=os.path.join(barcode_recovery_dir_merge, f"{run_name}_merge-stats.csv"),
        clean_headers_log=os.path.join(preprocessing_dir_merge, "logs/clean_headers/clean_headers.log"),
        merge_exon_remove=os.path.join(barcode_recovery_dir_merge, "logs/exonerate_int_cleanup_complete.txt"),
        # Concat mode inputs
        concat_summary_csv=os.path.join(barcode_recovery_dir_concat, f"{run_name}_concat-stats.csv"),
        concat_logs=os.path.join(preprocessing_dir_concat, "logs/concat/concat_reads.log"),
        trim_galore_logs=os.path.join(preprocessing_dir_concat, "logs/trim_galore/trim_galore.log"),
        concat_exon_remove=os.path.join(barcode_recovery_dir_concat, "logs/exonerate_int_cleanup_complete.txt"),
        # Wait for fasta cleanup to complete
        merge_fasta_cleanup=os.path.join(barcode_recovery_dir_merge, "logs/fasta_cleaner/fasta_cleaner_complete.txt"),
        concat_fasta_cleanup=os.path.join(barcode_recovery_dir_concat, "logs/fasta_cleaner/fasta_cleaner_complete.txt")
    output:
        merge_complete=touch(os.path.join(preprocessing_dir_merge, "logs/final_cleanup_complete.txt")),
        concat_complete=touch(os.path.join(preprocessing_dir_concat, "logs/final_cleanup_complete.txt"))
    threads: 1
    resources:
        mem_mb=3072
    retries: 2
    run:      
        # Function to clean up a specific mode's files
        def cleanup_mode_files(mode_dir, mode_name):
            removed_files = 0
            removed_dirs = []
            cleaned_dirs = []
            
            # Remove empty .log files in mge subdir - both modes
            mge_log_dir = os.path.join(mode_dir, "logs/mge")
            if os.path.exists(mge_log_dir):
                # Find all .log files in subdirectories of mge/
                mge_log_files = glob.glob(os.path.join(mge_log_dir, "**", "*.log"), recursive=True)
                removed_empty_logs = 0
                skipped_nonempty_logs = 0
                
                for mge_log_file in mge_log_files:
                    # Check if the log file is empty (0 bytes)
                    if os.path.getsize(mge_log_file) == 0:
                        print(f"[{mode_name}] Removing empty MGE log file: {mge_log_file}")
                        os.remove(mge_log_file)
                        removed_empty_logs += 1
                    else:
                        print(f"[{mode_name}] Keeping non-empty MGE log file: {mge_log_file} ({os.path.getsize(mge_log_file)} bytes)")
                        skipped_nonempty_logs += 1
                
                print(f"[{mode_name}] Removed {removed_empty_logs} empty MGE log files")
                print(f"[{mode_name}] Kept {skipped_nonempty_logs} non-empty MGE log files")
                removed_files += removed_empty_logs
                
                cleaned_dirs.append(mge_log_dir)
            else:
                print(f"[{mode_name}] MGE log directory not found: {mge_log_dir}")
            
            # Remove rename_complete.txt - both modes
            rename_complete_file = os.path.join(mode_dir, "logs/rename_consensus/rename_complete.txt")
            if os.path.exists(rename_complete_file):
                print(f"[{mode_name}] Removing rename_complete.txt file: {rename_complete_file}")
                os.remove(rename_complete_file)
                removed_files += 1
            else:
                print(f"[{mode_name}] rename_complete.txt file not found: {rename_complete_file}")
            
            # Return cleanup statistics
            return {
                "removed_files": removed_files,
                "removed_dirs": removed_dirs,
                "cleaned_dirs": cleaned_dirs
            }
        
        # Cleanup for merge mode (using barcode_recovery_dir for MGE logs)
        merge_stats = cleanup_mode_files(barcode_recovery_dir_merge, "merge")
        
        # Cleanup for concat mode (using barcode_recovery_dir for MGE logs)
        concat_stats = cleanup_mode_files(barcode_recovery_dir_concat, "concat")
        
        # Write completion message and cleanup info to merge output file
        with open(output.merge_complete, 'w') as f:
            f.write("Cleanup complete at " + datetime.now().strftime("%Y-%m-%d %H:%M:%S") + "\n\n")
            f.write(f"Total log files removed: {merge_stats['removed_files']}\n")            
            f.write("Directories fully removed:\n")
            if merge_stats['removed_dirs']:
                for dir_path in merge_stats['removed_dirs']:
                    f.write(f"- {dir_path}\n")
            else:
                f.write("- No directories were fully removed\n")
            
            f.write("\nDirectories cleaned of log files:\n")
            if merge_stats['cleaned_dirs']:
                for dir_path in merge_stats['cleaned_dirs']:
                    f.write(f"- {dir_path}\n")
            else:
                f.write("- No directories were cleaned\n")
            
            f.write("\nPreprocessing mode: merge")
        
        # Write completion message and cleanup info to concat output file
        with open(output.concat_complete, 'w') as f:
            f.write("Cleanup complete at " + datetime.now().strftime("%Y-%m-%d %H:%M:%S") + "\n\n")
            
            f.write(f"Total log files removed: {concat_stats['removed_files']}\n")
            
            f.write("Directories fully removed:\n")
            if concat_stats['removed_dirs']:
                for dir_path in concat_stats['removed_dirs']:
                    f.write(f"- {dir_path}\n")
            else:
                f.write("- No directories were fully removed\n")
            
            f.write("\nDirectories cleaned of log files:\n")
            if concat_stats['cleaned_dirs']:
                for dir_path in concat_stats['cleaned_dirs']:
                    f.write(f"- {dir_path}\n")
            else:
                f.write("- No directories were cleaned\n")
            
            f.write("\nPreprocessing mode: concat")


# ----- SE MitoGeneExtractor (mirrors MitoGeneExtractor_concat) -----
rule MitoGeneExtractor_se:
    input:
        DNA=get_mge_input_se,
        AA=get_protein_reference
    output:
        alignment=os.path.join(barcode_recovery_dir_se, "alignment/{sample}_r_{r}_s_{s}_align_{sample}.fas"),
        consensus=os.path.join(barcode_recovery_dir_se, "consensus/{sample}_r_{r}_s_{s}_con_{sample}.fas")
    log:
        out=os.path.join(barcode_recovery_dir_se, "out/{sample}_r_{r}_s_{s}_summary.out"),
        err=os.path.join(barcode_recovery_dir_se, "err/{sample}_r_{r}_s_{s}_summary.err")
    params:
        mge_executor=MGE_EXECUTABLE,
        n=n,
        C=C,
        t=t,
        output_dir=barcode_recovery_dir_se,
        vulgar_dir=lambda wildcards: os.path.join(barcode_recovery_dir_se, f"logs/mge/{wildcards.sample}_r_{wildcards.r}_s_{wildcards.s}/")
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["MitoGeneExtractor"]["mem_mb"] * attempt,
        slurm_partition=rule_resources["MitoGeneExtractor"]["partition"]
    retries: 4
    shell:
        """
        set -euo pipefail

        # Check if input file is empty or missing
        if [ ! -s "{input.DNA}" ]; then
            echo "===== EMPTY/MISSING INPUT DETECTED =====" >> {log.out}
            echo "Started: $(date)" >> {log.out}
            echo "Sample: {wildcards.sample}" >> {log.out}
            echo "Mode: se" >> {log.out}
            echo "Input file: {input.DNA}" >> {log.out}

            if [ ! -f "{input.DNA}" ]; then
                echo "Input file does not exist" >> {log.out}
            else
                FILE_SIZE=$(stat -c%s "{input.DNA}" 2>/dev/null || echo "unknown")
                echo "Input file size: $FILE_SIZE bytes" >> {log.out}
                echo "File is empty - likely all reads were filtered out during trimming" >> {log.out}
            fi

            echo "Creating mock outputs..." >> {log.out}
            echo ">Consensus__{wildcards.sample}" > {output.consensus}
            touch {output.alignment}
            touch {log.err}
            echo "Mock outputs created successfully: $(date)" >> {log.out}
            echo "=========================================" >> {log.out}
            exit 0
        fi

        mkdir -p {params.vulgar_dir}

        echo "===== Job Info =====" >> {log.out}
        echo "Started: $(date)" >> {log.out}
        echo "Node: $(hostname)" >> {log.out}
        echo "Memory allocated: $(({resources.mem_mb} / 1024))GB" >> {log.out}
        echo "Threads: {threads}" >> {log.out}
        echo "Mode: se" >> {log.out}
        echo "Sample: {wildcards.sample}" >> {log.out}

        FILE_SIZE_BYTES=$(stat -c%s "{input.DNA}")
        FILE_SIZE_GB=$(echo "scale=2; $FILE_SIZE_BYTES / 1024 / 1024 / 1024" | bc)
        NUM_SEQUENCES=$(grep -c "^@" "{input.DNA}" || echo "0")

        echo "===== Input Files =====" >> {log.out}
        echo "Input file: {input.DNA}" >> {log.out}
        echo "File size: ${{FILE_SIZE_GB}}GB ($FILE_SIZE_BYTES bytes)" >> {log.out}
        echo "Number of sequences: $NUM_SEQUENCES" >> {log.out}
        echo "Input AA file: {input.AA}" >> {log.out}

        echo "===== Starting MitoGeneExtractor =====" >> {log.out}

        time {params.mge_executor} \\
        -q {input.DNA} -p {input.AA} \\
        -o {params.output_dir}/alignment/{wildcards.sample}_r_{wildcards.r}_s_{wildcards.s}_align_ \\
        -c {params.output_dir}/consensus/{wildcards.sample}_r_{wildcards.r}_s_{wildcards.s}_con_ \\
        -V {params.vulgar_dir} \\
        -r {wildcards.r} -s {wildcards.s} \\
        -n {params.n} \\
        -C {params.C} \\
        -t {params.t} \\
        --temporaryDirectory {params.output_dir} \\
        --verbosity 10 \\
        >> {log.out} 2>> {log.err} || {{
            EXIT_CODE=$?
            echo "===== MitoGeneExtractor Failed =====" >> {log.out}
            echo "Exit code: $EXIT_CODE" >> {log.out}
            echo "Finished: $(date)" >> {log.out}
            if [ $EXIT_CODE -eq 137 ]; then
                echo "ERROR: Process killed (likely out of memory)" >> {log.out}
            elif [ $EXIT_CODE -eq 124 ]; then
                echo "ERROR: Process timed out" >> {log.out}
            else
                echo "ERROR: failure (exit code $EXIT_CODE)" >> {log.out}
            fi
            free -h >> {log.out}
            exit $EXIT_CODE
        }}

        echo "===== Job Completed Successfully =====" >> {log.out}
        echo "Finished: $(date)" >> {log.out}

        if [ -f "{output.consensus}" ] && [ -f "{output.alignment}" ]; then
            CONSENSUS_LINES=$(wc -l < {output.consensus})
            echo "===== Output Summary =====" >> {log.out}
            echo "Consensus lines: $CONSENSUS_LINES" >> {log.out}
            if [ $CONSENSUS_LINES -gt 1 ]; then
                echo "Consensus contains sequence data" >> {log.out}
            else
                echo "WARNING: Consensus may be header-only" >> {log.out}
            fi
        else
            echo "ERROR: Expected output files not found" >> {log.out}
            exit 1
        fi
        """

# ----- SE rename + combine consensus (mirrors rename_and_combine_cons_concat) -----
rule rename_and_combine_cons_se:
    input:
        consensus_files=get_all_mge_consensus_files("se")
    output:
        complete=os.path.join(barcode_recovery_dir_se, "logs/rename_consensus/rename_complete.txt"),
        concat_cons=os.path.join(barcode_recovery_dir_se, f"consensus/{run_name}_cons_combined-se.fasta")
    log:
        os.path.join(barcode_recovery_dir_se, "logs/rename_consensus/rename_fasta.log")
    params:
        script=os.path.join(workflow.basedir, "scripts/rename_headers.py"),
        preprocessing_mode="se"
    threads: rule_resources["rename_and_combine_cons"]["threads"]
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["rename_and_combine_cons"]["mem_mb"] * attempt
    retries: 3
    shell:
        """
        set -euo pipefail

        for file in {input.consensus_files}; do
            if [ ! -f "$file" ]; then
                echo "Error: Required input file does not exist: $file" >> {log}
                exit 1
            fi
        done

        python {params.script} \\
            --input-files {input.consensus_files} \\
            --concatenated-consensus {output.concat_cons} \\
            --complete-file={output.complete} \\
            --log-file={log} \\
            --preprocessing-mode={params.preprocessing_mode} \\
            --threads={threads}

        touch {output.complete}
        """

# ----- SE exonerate intermediate cleanup (se dir only) -----
rule remove_exonerate_intermediates_se:
    input:
        se_consensus=os.path.join(barcode_recovery_dir_se, f"consensus/{run_name}_cons_combined-se.fasta"),
        se_rename_complete=os.path.join(barcode_recovery_dir_se, "logs/rename_consensus/rename_complete.txt")
    output:
        se_completion=os.path.join(barcode_recovery_dir_se, "logs/exonerate_int_cleanup_complete.txt")
    threads: 1
    resources:
        mem_mb=2048
    retries: 2
    run:
        with open(output.se_completion, 'w') as log:
            log.write(f"Starting exonerate intermediate file cleanup for se mode: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            log.write(f"Mode directory: {barcode_recovery_dir_se}\n\n")
            exonerate_pattern = os.path.join(barcode_recovery_dir_se, "Concatenated_exonerate_input_*")
            exonerate_files = glob.glob(exonerate_pattern)
            removed_files = 0
            total_size_saved = 0
            if not exonerate_files:
                log.write("No Concatenated_exonerate_input_* files found\n")
            else:
                log.write(f"Found {len(exonerate_files)} files to remove:\n")
                for file_path in exonerate_files:
                    try:
                        file_size = os.path.getsize(file_path)
                        os.remove(file_path)
                        log.write(f"  - Removed: {os.path.basename(file_path)} ({file_size / (1024*1024):.2f} MB)\n")
                        removed_files += 1
                        total_size_saved += file_size
                    except OSError as e:
                        log.write(f"  - ERROR removing {file_path}: {e}\n")
            log.write(f"\nCLEANUP SUMMARY for se mode:\n")
            log.write(f"  Total files removed: {removed_files}\n")
            log.write(f"  Total disk space saved: {total_size_saved / (1024*1024):.2f} MB\n")
            log.write(f"  Completed: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")

# ----- SE alignment log (mirrors create_alignment_log_concat) -----
rule create_alignment_log_se:
    input:
        concat_cons=os.path.join(barcode_recovery_dir_se, f"consensus/{run_name}_cons_combined-se.fasta")
    output:
        alignment_log=os.path.join(barcode_recovery_dir_se, "logs/mge/alignment_files.log")
    params:
        alignment_dir=os.path.join(barcode_recovery_dir_se, "alignment")
    threads: 1
    resources:
        mem_mb=4096
    retries: 3
    shell:
        "find {params.alignment_dir} -name '*_align_*.fas' -type f 2>/dev/null | sort > {output.alignment_log} || touch {output.alignment_log}"

# ===== FASTA CLEANER RULES - SE MODE =====
rule human_cox1_filter_se:
    input:
        alignment_log=os.path.join(barcode_recovery_dir_se, "logs/mge/alignment_files.log")
    output:
        filtered_files=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/01_human_filtered/human_filtered.txt"),
        metrics=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/01_human_filtered/human_filter_metrics.csv")
    params:
        script=os.path.join(workflow.basedir, "scripts/01_human_cox1_filter.py"),
        output_dir=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/01_human_filtered"),
        human_threshold=fasta_cleaner_params["human_threshold"]
    log:
        os.path.join(barcode_recovery_dir_se, "logs/fasta_cleaner/01_human_filter.log")
    threads: rule_resources["human_cox1_filter"]["threads"]
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["human_cox1_filter"]["mem_mb"] * attempt
    retries: 2
    shell:
        """
        set -euo pipefail
        mkdir -p {params.output_dir}
        python {params.script} \\
            --input-log {input.alignment_log} \\
            --output-dir {params.output_dir} \\
            --filtered-files-list {output.filtered_files} \\
            --metrics-csv {output.metrics} \\
            --human-threshold {params.human_threshold} \\
            --threads {threads} \\
            2>&1 | tee {log}
        """

rule at_content_filter_se:
    input:
        human_filtered=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/01_human_filtered/human_filtered.txt")
    output:
        filtered_files=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/02_at_filtered/at_filtered.txt"),
        summary_metrics=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/02_at_filtered/at_filter_summary.csv")
    params:
        script=os.path.join(workflow.basedir, "scripts/02_at_content_filter.py"),
        output_dir=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/02_at_filtered"),
        at_threshold=fasta_cleaner_params["at_difference"],
        at_mode=fasta_cleaner_params["at_mode"],
        consensus_threshold=fasta_cleaner_params["consensus_threshold"]
    log:
        os.path.join(barcode_recovery_dir_se, "logs/fasta_cleaner/02_at_filter.log")
    threads: rule_resources["at_content_filter"]["threads"]
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["at_content_filter"]["mem_mb"] * attempt
    retries: 2
    shell:
        """
        set -euo pipefail
        mkdir -p {params.output_dir}
        python -u {params.script} \\
            --input_files {input.human_filtered} \\
            --output_dir {params.output_dir} \\
            --filtered-files-list {output.filtered_files} \\
            --at_threshold {params.at_threshold} \\
            --at_mode {params.at_mode} \\
            --consensus_threshold {params.consensus_threshold} \\
            --threads {threads} \\
            2>&1 | tee {log}
        """

rule statistical_outlier_filter_se:
    input:
        filtered_files=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/02_at_filtered/at_filtered.txt")
    output:
        filtered_files=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/03_outlier_filtered/outlier_filtered.txt"),
        summary_metrics=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/03_outlier_filtered/outlier_filter_summary_metrics.csv"),
        individual_metrics=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/03_outlier_filtered/outlier_filter_individual_metrics.csv")
    params:
        script=os.path.join(workflow.basedir, "scripts/03_statistical_outlier_filter.py"),
        output_dir=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/03_outlier_filtered"),
        outlier_percentile=fasta_cleaner_params["outlier_percentile"],
        consensus_threshold=fasta_cleaner_params["consensus_threshold"]
    log:
        os.path.join(barcode_recovery_dir_se, "logs/fasta_cleaner/03_outlier_filter.log")
    threads: rule_resources["statistical_outlier_filter"]["threads"]
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["statistical_outlier_filter"]["mem_mb"] * attempt
    retries: 2
    shell:
        """
        set -euo pipefail
        mkdir -p {params.output_dir}
        python {params.script} \\
            --input-files-list {input.filtered_files} \\
            --output-dir {params.output_dir} \\
            --filtered-files-list {output.filtered_files} \\
            --summary-csv {output.individual_metrics} \\
            --metrics-csv {output.summary_metrics} \\
            --outlier-percentile {params.outlier_percentile} \\
            --consensus-threshold {params.consensus_threshold} \\
            --threads {threads} \\
            2>&1 | tee {log}
        """

if use_reference_filtering:
    rule reference_filter_se:
        input:
            filtered_files=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/03_outlier_filtered/outlier_filtered.txt")
        output:
            filtered_files=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/04_reference_filtered/reference_filtered.txt"),
            metrics=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/04_reference_filtered/reference_filter_metrics.csv")
        params:
            script=os.path.join(workflow.basedir, "scripts/04_reference_filter.py"),
            output_dir=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/04_reference_filtered"),
            reference_dir=reference_dir,
            filter_mode=fasta_cleaner_params["reference_filter_mode"],
        log:
            os.path.join(barcode_recovery_dir_se, "logs/fasta_cleaner/04_reference_filter.log")
        threads: rule_resources["reference_filter"]["threads"]
        resources:
            mem_mb=lambda wildcards, attempt: rule_resources["reference_filter"]["mem_mb"] * attempt
        retries: 2
        shell:
            """
            set -euo pipefail
            mkdir -p {params.output_dir}
            python {params.script} \\
                --input-files-list {input.filtered_files} \\
                --output-dir {params.output_dir} \\
                --filtered-files-list {output.filtered_files} \\
                --metrics-csv {output.metrics} \\
                --reference-dir {params.reference_dir} \\
                --filter-mode {params.filter_mode} \\
                --threads {threads} \\
                2>&1 | tee {log}
            """

    rule consensus_generation_se:
        input:
            filtered_files=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/04_reference_filtered/reference_filtered.txt")
        output:
            concat_consensus=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/cleaned_cons_combined.fasta"),
            consensus_metrics=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/05_cleaned_consensus/cleaned_cons_metrics-se.csv")
        params:
            script=os.path.join(workflow.basedir, "scripts/05_consensus_generator.py"),
            output_dir=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/05_cleaned_consensus"),
            consensus_threshold=fasta_cleaner_params["consensus_threshold"],
            preprocessing_mode="se"
        log:
            os.path.join(barcode_recovery_dir_se, "logs/fasta_cleaner/05_consensus_generation.log")
        threads: rule_resources["consensus_generation"]["threads"]
        resources:
            mem_mb=lambda wildcards, attempt: rule_resources["consensus_generation"]["mem_mb"] * attempt
        retries: 2
        shell:
            """
            set -euo pipefail
            python {params.script} \\
                --input-files-list {input.filtered_files} \\
                --output-dir {params.output_dir} \\
                --consensus-fasta {output.concat_consensus} \\
                --consensus-metrics {output.consensus_metrics} \\
                --consensus-threshold {params.consensus_threshold} \\
                --preprocessing-mode {params.preprocessing_mode} \\
                --threads {threads} \\
                2>&1 | tee {log}
            """

    rule aggregate_filter_metrics_se:
        input:
            human_metrics=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/01_human_filtered/human_filter_metrics.csv"),
            at_metrics=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/02_at_filtered/at_filter_summary.csv"),
            outlier_metrics=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/03_outlier_filtered/outlier_filter_individual_metrics.csv"),
            reference_metrics=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/04_reference_filtered/reference_filter_metrics.csv"),
            consensus_metrics=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/05_cleaned_consensus/cleaned_cons_metrics-se.csv")
        output:
            completion_flag=os.path.join(barcode_recovery_dir_se, "logs/fasta_cleaner_complete.txt"),
            combined_statistics=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/combined_statistics.csv")
        params:
            script=os.path.join(workflow.basedir, "scripts/06_aggregate_filter_metrics.py"),
            output_dir=os.path.join(barcode_recovery_dir_se, "fasta_cleaner")
        log:
            os.path.join(barcode_recovery_dir_se, "logs/fasta_cleaner/06_aggregate_metrics.log")
        threads: 2
        resources:
            mem_mb=3076
        retries: 2
        shell:
            """
            set -euo pipefail
            python {params.script} \\
                --human-metrics {input.human_metrics} \\
                --at-metrics {input.at_metrics} \\
                --outlier-metrics {input.outlier_metrics} \\
                --reference-metrics {input.reference_metrics} \\
                --consensus-metrics {input.consensus_metrics} \\
                --output-dir {params.output_dir} \\
                --combined-statistics {output.combined_statistics} \\
                --threads {threads} \\
                2>&1 | tee {log}
            touch {output.completion_flag}
            """

else:
    rule consensus_generation_se:
        input:
            filtered_files=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/03_outlier_filtered/outlier_filtered.txt")
        output:
            concat_consensus=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/cleaned_cons_combined.fasta"),
            consensus_metrics=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/05_cleaned_consensus/cleaned_cons_metrics-se.csv")
        params:
            script=os.path.join(workflow.basedir, "scripts/05_consensus_generator.py"),
            output_dir=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/05_cleaned_consensus"),
            consensus_threshold=fasta_cleaner_params["consensus_threshold"],
            preprocessing_mode="se"
        log:
            os.path.join(barcode_recovery_dir_se, "logs/fasta_cleaner/05_consensus_generation.log")
        threads: rule_resources["consensus_generation"]["threads"]
        resources:
            mem_mb=lambda wildcards, attempt: rule_resources["consensus_generation"]["mem_mb"] * attempt
        retries: 2
        shell:
            """
            set -euo pipefail
            python {params.script} \\
                --input-files-list {input.filtered_files} \\
                --output-dir {params.output_dir} \\
                --consensus-fasta {output.concat_consensus} \\
                --consensus-metrics {output.consensus_metrics} \\
                --consensus-threshold {params.consensus_threshold} \\
                --preprocessing-mode {params.preprocessing_mode} \\
                --threads {threads} \\
                2>&1 | tee {log}
            """

    rule aggregate_filter_metrics_se:
        input:
            human_metrics=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/01_human_filtered/human_filter_metrics.csv"),
            at_metrics=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/02_at_filtered/at_filter_summary.csv"),
            outlier_metrics=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/03_outlier_filtered/outlier_filter_individual_metrics.csv"),
            consensus_metrics=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/05_cleaned_consensus/cleaned_cons_metrics-se.csv")
        output:
            completion_flag=os.path.join(barcode_recovery_dir_se, "logs/fasta_cleaner_complete.txt"),
            combined_statistics=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/combined_statistics.csv")
        params:
            script=os.path.join(workflow.basedir, "scripts/06_aggregate_filter_metrics.py"),
            output_dir=os.path.join(barcode_recovery_dir_se, "fasta_cleaner")
        log:
            os.path.join(barcode_recovery_dir_se, "logs/fasta_cleaner/06_aggregate_metrics.log")
        threads: 1
        resources:
            mem_mb=3076
        retries: 2
        shell:
            """
            set -euo pipefail
            python {params.script} \\
                --human-metrics {input.human_metrics} \\
                --at-metrics {input.at_metrics} \\
                --outlier-metrics {input.outlier_metrics} \\
                --consensus-metrics {input.consensus_metrics} \\
                --output-dir {params.output_dir} \\
                --combined-statistics {output.combined_statistics} \\
                --threads {threads} \\
                2>&1 | tee {log}
            touch {output.completion_flag}
            """

# ----- SE intermediate fasta cleanup (se dir only) -----
rule remove_fasta_cleaner_files_se:
    input:
        se_cleaning_complete=os.path.join(barcode_recovery_dir_se, "logs/fasta_cleaner_complete.txt"),
        se_consensus=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/cleaned_cons_combined.fasta")
    output:
        se_cleanup_complete=os.path.join(barcode_recovery_dir_se, "logs/fasta_cleaner/fasta_cleaner_complete.txt")
    params:
        se_dir=barcode_recovery_dir_se,
        use_ref_filtering=str(use_reference_filtering)
    threads: 1
    resources:
        mem_mb=2048
    retries: 2
    run:
        removed_files = 0
        total_size_saved = 0
        with open(output.se_cleanup_complete, 'w') as log:
            log.write(f"Starting intermediate fasta cleanup for se mode: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            log.write(f"Mode directory: {params.se_dir}\n")
            log.write(f"Reference filtering enabled: {params.use_ref_filtering}\n\n")
            filter_dirs = ["01_human_filtered", "02_at_filtered", "03_outlier_filtered"]
            if params.use_ref_filtering == "True":
                filter_dirs.append("04_reference_filtered")
            for filter_dir in filter_dirs:
                dir_path = os.path.join(params.se_dir, "fasta_cleaner", filter_dir)
                if not os.path.exists(dir_path):
                    log.write(f"[{filter_dir}] Directory not found: {dir_path}\n")
                    continue
                fasta_files = []
                for pattern in [os.path.join(dir_path, "*.fasta"), os.path.join(dir_path, "*.fas"),
                                os.path.join(dir_path, "**", "*.fasta"), os.path.join(dir_path, "**", "*.fas")]:
                    fasta_files.extend(glob.glob(pattern, recursive=True))
                fasta_files = list(set(fasta_files))
                for fasta_file in fasta_files:
                    try:
                        file_size = os.path.getsize(fasta_file)
                        os.remove(fasta_file)
                        removed_files += 1
                        total_size_saved += file_size
                    except OSError as e:
                        log.write(f"  - ERROR removing {fasta_file}: {e}\n")
                log.write(f"[{filter_dir}] cleaned\n")
            log.write(f"\nCLEANUP SUMMARY for se mode:\n")
            log.write(f"  Total files removed: {removed_files}\n")
            log.write(f"  Total disk space saved: {total_size_saved / (1024*1024):.2f} MB\n")
            log.write(f"  Completed: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")

# ----- SE stats compilation (mirrors extract_stats_to_csv_concat) -----
rule extract_stats_to_csv_se:
    input:
        alignment_log=os.path.join(barcode_recovery_dir_se, "logs/mge/alignment_files.log"),
        out_files=expand(
            os.path.join(barcode_recovery_dir_se, "out/{sample}_r_{r}_s_{s}_summary.out"),
            sample=list(samples.keys()),
            r=r,
            s=s
        ),
        pre_fasta=os.path.join(barcode_recovery_dir_se, f"consensus/{run_name}_cons_combined-se.fasta"),
        cleaning_csv=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/combined_statistics.csv"),
        cleaner_complete=os.path.join(barcode_recovery_dir_se, "logs/fasta_cleaner_complete.txt"),
        post_fasta=os.path.join(barcode_recovery_dir_se, "fasta_cleaner/cleaned_cons_combined.fasta"),
        ref_seqs=os.path.join(main_output_dir, "02_references/sequence_references.csv") if run_gene_fetch else []
    output:
        stats=os.path.join(barcode_recovery_dir_se, f"{run_name}_se-stats.csv")
    log:
        os.path.join(barcode_recovery_dir_se, "logs/compile_barcoding_stats.log")
    params:
        script=os.path.join(workflow.basedir, "scripts/compile_barcoding_stats.py"),
        out_dir=os.path.join(barcode_recovery_dir_se, "out/"),
        ref_seqs_arg=lambda wildcards, input: f"--ref_seqs {input.ref_seqs}" if run_gene_fetch else ""
    resources:
        mem_mb=lambda wildcards, attempt: rule_resources["extract_stats_to_csv"]["mem_mb"] * attempt
    retries: 3
    shell:
        """
        set -euo pipefail
        python {params.script} -a {input.alignment_log} -o {output.stats} -od {params.out_dir} -c {input.cleaning_csv} {params.ref_seqs_arg} --pre-fasta {input.pre_fasta} --post-fasta {input.post_fasta} --log-file {log}
        """
