"""
Leech Snakemake Pipeline
========================

Trains models for pairwise comparisons, including:
- Charged vs uncharged tRNA classification
- Amino acid pairwise classification

All comparisons are defined uniformly via TSV spec files (see config/comparisons_*.tsv).

Supports both SLURM and LSF with GPU acceleration.

Usage:
    # Dry run
    snakemake --profile profiles/slurm -n

    # Run with SLURM
    snakemake --profile profiles/slurm

    # Run with LSF
    snakemake --profile profiles/lsf

    # Run locally (for testing)
    snakemake --cores 8
"""

from pathlib import Path


# Load configuration
configfile: str(Path(workflow.basedir).parent / "config" / "config.yaml")


# Map pod5 -> raw_pod5 for dorado workflow compatibility
# Samples should be provided via --configfile (e.g., config/samples-alpine.yaml)
if "samples" in config:
    for sample_name, sample_data in config["samples"].items():
        if "pod5" in sample_data and "raw_pod5" not in sample_data:
            sample_data["raw_pod5"] = sample_data.pop("pod5")


# Include common variables and functions
include: "rules/common.smk"
# Include modular rule files
include: "rules/documentation.smk"
include: "rules/dorado.smk"
include: "rules/inspect_pod5.smk"
include: "rules/prepare.smk"
include: "rules/grid_search.smk"
include: "rules/train.smk"
include: "rules/inference.smk"
include: "rules/evaluate.smk"
include: "rules/compare_models.smk"


# ============================================================================
# Target rules
# ============================================================================


rule all:
    """Default target: run all analyses."""
    input:
        # Run summary document (generated first)
        RUN_SUMMARY_FILE,
        # Choose based on compare_models setting
        METRICS_DIR + "/comparison/aggregate/pairwise_architecture_comparison.tsv.gz"
        if COMPARE_MODELS and AA_PAIRS
        else (METRICS_DIR + "/pairwise_summary.tsv.gz" if AA_PAIRS else []),


# Single model mode (compare_models: false)
rule all_single_model:
    """Run single model analysis (default model from config)."""
    input:
        # Run summary document (generated first)
        RUN_SUMMARY_FILE,
        # Pairwise classifiers (includes charged_uncharged if configured)
        expand(
            METRICS_DIR + "/pairwise/{pair}/test_metrics.json",
            pair=AA_PAIRS,
        )
        if AA_PAIRS
        else [],
        # Summary reports
        METRICS_DIR + "/pairwise_summary.tsv.gz" if AA_PAIRS else [],


# Multi-architecture comparison mode (compare_models: true)
rule all_compare_models:
    """Run multi-architecture comparison."""
    input:
        # Run summary document (generated first)
        RUN_SUMMARY_FILE,
        # Architecture comparison results (includes charged_uncharged if configured)
        METRICS_DIR + "/comparison/aggregate/pairwise_architecture_comparison.tsv.gz"
        if AA_PAIRS
        else [],
        METRICS_DIR + "/comparison/aggregate/pairwise_architecture_summary.txt"
        if AA_PAIRS
        else [],


rule all_prepare:
    """Prepare all training chunks (without splitting)."""
    input:
        # Run summary document (generated first)
        RUN_SUMMARY_FILE,
        expand(
            CHUNKS_DIR + "/{sample}/all.npz",
            sample=SAMPLES,
        ),


rule all_merge:
    """Merge all prepared chunks and split at read level."""
    input:
        # Run summary document (generated first)
        RUN_SUMMARY_FILE,
        expand(CHUNKS_DIR + "/merged/pairwise/{pair}/train.npz", pair=AA_PAIRS)
        if AA_PAIRS
        else [],
        expand(CHUNKS_DIR + "/merged/pairwise/{pair}/val.npz", pair=AA_PAIRS)
        if AA_PAIRS
        else [],
        expand(CHUNKS_DIR + "/merged/pairwise/{pair}/test.npz", pair=AA_PAIRS)
        if AA_PAIRS
        else [],


rule all_grid_search:
    """Run all grid searches (single model mode)."""
    input:
        expand(
            MODELS_DIR + "/grid_search/pairwise/{pair}/best_params.json", pair=AA_PAIRS
        )
        if AA_PAIRS
        else [],


rule all_grid_search_comparison:
    """Run grid searches for all architectures (comparison mode)."""
    input:
        expand(
            MODELS_DIR + "/grid_search/pairwise/{pair}/{arch}/best_params.json",
            pair=AA_PAIRS,
            arch=MODEL_ARCHITECTURES,
        )
        if AA_PAIRS
        else [],


rule all_train:
    """Train all models (single model mode)."""
    input:
        expand(MODELS_DIR + "/pairwise/{pair}/model_best.pt", pair=AA_PAIRS)
        if AA_PAIRS
        else [],


rule all_train_comparison:
    """Train all architectures (comparison mode)."""
    input:
        expand(
            MODELS_DIR + "/comparison/pairwise/{pair}/{arch}/model_best.pt",
            pair=AA_PAIRS,
            arch=MODEL_ARCHITECTURES,
        )
        if AA_PAIRS
        else [],


rule all_infer:
    """Run all inference."""
    input:
        expand(
            INFER_DIR + "/pairwise/{pair}/{sample}_predictions.bam",
            pair=AA_PAIRS,
            sample=SAMPLES,
        )
        if AA_PAIRS
        else [],


rule all_merge_pods:
    """Merge all raw POD5 files."""
    input:
        expand(
            get_project_path(config.get("pod5_dir", "results/pod5"))
            + "/{sample}/{sample}.pod5",
            sample=SAMPLES,
        ),


rule all_rebasecall:
    """Rebasecall all samples with dorado."""
    input:
        expand(
            get_project_path(config.get("rebasecall_dir", "results/bam/rebasecall"))
            + "/{sample}/{sample}.rbc.bam",
            sample=SAMPLES,
        ),


rule all_align:
    """Align all rebasecalled samples to reference."""
    input:
        expand(
            get_project_path(config.get("rebasecall_dir", "results/bam/rebasecall"))
            + "/{sample}/{sample}.aligned.bam",
            sample=SAMPLES,
        ),


rule all_inspect_pod5:
    """Inspect all raw POD5 files for corruption/errors."""
    input:
        "results/pod5_inspection/pod5_inspection_master_report.tsv.gz",
        "results/pod5_inspection/pod5_inspection_summary.txt",
