###############################################################################
# Snakefile: Run IQ-TREE per genus
#
# This workflow:
# 1. Reads the genome summary CSV
# 2. Extracts the list of genera to process
# 3. Uses GToTree outputs (alignment + partitions)
# 4. Runs IQ-TREE separately for each genus
# 5. Creates a ".done" file to mark successful completion
#
# ASSUMPTIONS:
# - GToTree has already been run for each genus
# - The following files exist for each genus:
#     • Aligned_SCGs.faa
#     • run_files/Partitions.txt
# - All paths are defined in config.yaml
###############################################################################

import pandas as pd
import os

# ---------------------------------------------------------------------------
# Load configuration file
# ---------------------------------------------------------------------------
# Snakemake automatically reads config.yaml and exposes values via `config`
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 file

# ---------------------------------------------------------------------------
# Load genome summary to determine which genera to process
# ---------------------------------------------------------------------------
# Each unique genus will be processed independently by IQ-TREE
genomes_df = pd.read_csv(SUMMARY_FILE)
GENERA = list(genomes_df["Genus"].unique())

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

#############################################
# Rule: run_iqtree
#############################################
# Runs IQ-TREE for a single genus
#
# Inputs:
# - Concatenated SCG alignment produced by GToTree
# - Partition file describing gene boundaries
#
# Output:
# - A ".done" file indicating successful IQ-TREE completion
#############################################
rule run_iqtree:
    input:
        # Alignment of single-copy genes from GToTree
        alignment=lambda w: os.path.join(
            RESULTS_DIR,
            f"{w.genus}_GToTree",
            "Aligned_SCGs.faa"
        ),

        # Partition file generated by GToTree
        partitions=lambda w: os.path.join(
            RESULTS_DIR,
            f"{w.genus}_GToTree",
            "run_files",
            "Partitions.txt"
        ),
    output:
        # Marker file to signal that IQ-TREE finished successfully
        touch(os.path.join(
            RESULTS_DIR,
            "{genus}_iqtree",
            "{genus}_iqtree.done"
        ))
    params:
        # Output directory for IQ-TREE results
        outdir=lambda w: os.path.join(
            RESULTS_DIR,
            f"{w.genus}_iqtree"
        ),

        # Prefix used by IQ-TREE for all output files
        prefix=lambda w: os.path.join(
            RESULTS_DIR,
            f"{w.genus}_iqtree",
            f"{w.genus}_iqtree_out"
        )
    conda:
        # Conda environment containing IQ-TREE
        "./envs/iqtree.yaml"
    shell:
        r"""
        # Create output directory if it does not exist
        mkdir -p {params.outdir}
        cd {params.outdir}

        # Run IQ-TREE with:
        # - ModelFinder Plus (MFP)
        # - 1000 ultrafast bootstrap replicates
        # - 4 CPU threads
        iqtree \
          -s {input.alignment} \
          -spp {input.partitions} \
          -m MFP \
          -bb 1000 \
          -nt 4 \
          -pre {params.prefix}
        """

