###############################################################################
# Snakefile: Run GToTree per genus
#
# This workflow:
# 1. Reads the genome summary CSV
# 2. Maps each genus to its corresponding HMM file
# 3. Runs GToTree independently for each genus
# 4. Uses a ".done" file to mark successful completion
#
# ASSUMPTIONS:
# - The previous Snakefile has already:
#   • Grouped genomes by genus
#   • Downloaded outgroups
#   • Created *_fasta_files.txt and *_final_GCF.txt per genus
# - All required paths are defined in config.yaml
###############################################################################

import pandas as pd
import os

# ---------------------------------------------------------------------------
# Load configuration file
# ---------------------------------------------------------------------------
# Snakemake reads config.yaml and exposes values via the `config` dictionary
configfile: "config.yaml"

# ---------------------------------------------------------------------------
# Read required paths from config.yaml
# ---------------------------------------------------------------------------
RESULTS_DIR = config["results_dir"]     # Base results directory
SUMMARY_FILE = config["genome_summary"] # Genome summary CSV

# ---------------------------------------------------------------------------
# Load genome summary file
# ---------------------------------------------------------------------------
# This file must contain at least:
# - Genus
# - HMM (path to the HMM file used by GToTree)
genomes_df = pd.read_csv(SUMMARY_FILE)

# A missing/blank Outgroup cell reads as NaN, which is truthy in Python
# (bool(float('nan')) is True) -- normalize to a real empty string first,
# same fix as pre_gtotree_snakefile, or the skip-check below silently does
# nothing and every genus (including ones with no outgroup) gets included.
if "Outgroup" in genomes_df.columns:
    genomes_df["Outgroup"] = (
        genomes_df["Outgroup"].fillna("").astype(str).str.strip()
        .replace({"nan": "", "None": "", "NaN": "", "<NA>": ""})
    )

# ---------------------------------------------------------------------------
# Build a mapping of genus → HMM
# ---------------------------------------------------------------------------
# One HMM is expected per genus
hmm_map = (
    genomes_df
    .drop_duplicates(subset=["Genus"])
    .set_index("Genus")[next(c for c in ('HMM', 'HMM_SCG') if c in genomes_df.columns)]
    .to_dict()
)

# ---------------------------------------------------------------------------
# Extract the list of genera to process
# ---------------------------------------------------------------------------
# Each genus will be processed independently by GToTree. Skip any genus with
# no resolved Outgroup -- the previous step (pre_gtotree_snakefile) already
# skips these and never creates their {genus}_fasta_files.txt /
# {genus}_final_GCF.txt inputs, so requesting a tree for them here is
# guaranteed to fail with MissingInputException. Without this check, a
# single genus with no outgroup takes down every OTHER genus's tree too,
# since Snakemake's default scheduler halts new job scheduling after any
# single job fails -- the same cascading-failure pattern fixed earlier in
# pre_gtotree_snakefile.
if "Outgroup" in genomes_df.columns:
    outgroup_map = (
        genomes_df
        .drop_duplicates(subset=["Genus"])
        .set_index("Genus")["Outgroup"]
        .to_dict()
    )
    GENERA = [g for g in hmm_map.keys() if outgroup_map.get(g)]
    skipped = [g for g in hmm_map.keys() if not outgroup_map.get(g)]
    if skipped:
        print(
            f"[NOSE] WARNING: No outgroup resolved for {len(skipped)} genus/genera "
            f"-- skipping GToTree for them (a tree can't be rooted without one): "
            f"{', '.join(skipped)}."
        )
else:
    GENERA = list(hmm_map.keys())

#############################################
# Rule: all
#############################################
# Entry point for the workflow
# Ensures GToTree has successfully completed for every genus
rule all:
    input:
        expand(
            os.path.join(RESULTS_DIR, "{genus}", "{genus}_GToTree.done"),
            genus=GENERA
        )

#############################################
# Rule: run_gtotree
#############################################
# Runs GToTree for a single genus
#
# Inputs:
# - A text file listing all FASTA/FNA files for the genus
# - A text file listing GCF accessions for reference genomes
#
# Output:
# - A ".done" file used as a completion marker
#############################################
rule run_gtotree:
    input:
        # List of genome FASTA/FNA files for this genus
        fasta_list=lambda w: os.path.join(
            RESULTS_DIR,
            w.genus,
            f"{w.genus}_fasta_files.txt"
        ),

        # List of GCF accessions for this genus
        accessions=lambda w: os.path.join(
            RESULTS_DIR,
            w.genus,
            f"{w.genus}_final_GCF.txt"
        ),
    output:
        # Marker file indicating successful GToTree execution
        touch(os.path.join(
            RESULTS_DIR,
            "{genus}",
            "{genus}_GToTree.done"
        ))
    params:
        # HMM file specific to this genus
        hmm=lambda w: hmm_map[w.genus],

        # Output directory for GToTree results
        outdir=lambda w: os.path.join(
            RESULTS_DIR,
            f"{w.genus}_GToTree"
        )
    conda:
        # Conda environment containing GToTree and dependencies
        "./envs/gtotree.yaml"
    shell:
        r"""
        # Move into the genus-specific results directory
        cd {RESULTS_DIR}/{wildcards.genus}

        # Run GToTree
        GToTree \
          -a {input.accessions} \
          -f {input.fasta_list} \
          -H {params.hmm} \
          -j 4 \
          -o {params.outdir}
        """

