# Main entrypoint of the workflow.
# Please follow the best practices:
# https://snakemake.readthedocs.io/en/stable/snakefiles/best_practices.html,
# in particular regarding the standardized folder structure mentioned there.


# load rules
# -----------------------------------------------------
# common.smk: Defines all variables + output files
include: "rules/common.smk"
include: "rules/io.smk"
include: "rules/tissue_id.smk"
include: "rules/histology.smk"
include: "rules/segmentation.smk"
include: "rules/cell_featurization.smk"
include: "rules/classification.smk"
include: "rules/plot.smk"


# optional messages, log and error handling
# -----------------------------------------------------
onstart:
    print("\n--- Analysis started ---\n")


onsuccess:
    print("\n--- Workflow finished! ---\n")


onerror:
    print("\n--- An error occurred! ---\n")


# target rules
# -----------------------------------------------------
rule all:
    input:
        # SDATA objects
        expand(
            "{output_dir}/done/to_sdata/{sample}.sdata.out",
            output_dir=config["GENERAL"]["output_dir"],
            sample=SAMPLES,
        ),
        # Segmentation
        (
            expand(
                "{output_dir}/done/segmentation/{sample}.{segmentation_key}.{table_suffix}",
                output_dir=config["GENERAL"]["output_dir"],
                sample=SAMPLES,
                segmentation_key=SEGMENTATION_KEYS,
                table_suffix=config.get("SEGMENTATION", {}).get("tables_suffix"),
            )
            if config.get("SEGMENTATION") is not None
            else []
        ),
        # Feature processing
        (
            expand(
                "{output_dir}/plots/qc/histogram/{segmentation_key}/{sample}.{segmentation_key}.{table_suffix}.qc.png",
                output_dir=config["GENERAL"]["output_dir"],
                sample=SAMPLES,
                segmentation_key=SEGMENTATION_KEYS,
                table_suffix=config.get("SEGMENTATION", {}).get("tables_suffix"),
            )
            if (config.get("FLUORESCENCE_PROCESSING") is not None)
            and (config.get("SEGMENTATION") is not None)
            else []
        ),
        (
            expand(
                "{output_dir}/plots/qc/overlay/{segmentation_key}/{sample}.{segmentation_key}.{table_suffix}.qc-overlay.png",
                output_dir=config["GENERAL"]["output_dir"],
                sample=SAMPLES,
                segmentation_key=SEGMENTATION_KEYS,
                table_suffix=config.get("SEGMENTATION", {}).get("tables_suffix"),
            )
            if (config.get("FLUORESCENCE_PROCESSING") is not None)
            and (config.get("SEGMENTATION") is not None)
            else []
        ),
        # GMM classification
        (
            expand(
                "{output_dir}/done/gmm_classification/{sample}.{segmentation_key}.{table_suffix}",
                output_dir=config["GENERAL"]["output_dir"],
                sample=SAMPLES,
                segmentation_key=SEGMENTATION_KEYS,
                table_suffix=config.get("SEGMENTATION", {}).get("tables_suffix"),
            )
            if (config.get("CLASSIFICATION") is not None)
            and (config.get("SEGMENTATION") is not None)
            else []
        ),
        (
            expand(
                "{output_dir}/plots/classification/histogram/{segmentation_key}/{sample}.{segmentation_key}.{table_suffix}.gmm.png",
                output_dir=config["GENERAL"]["output_dir"],
                sample=SAMPLES,
                segmentation_key=SEGMENTATION_KEYS,
                table_suffix=config.get("SEGMENTATION", {}).get("tables_suffix"),
            )
            if (config.get("CLASSIFICATION") is not None)
            and (config.get("SEGMENTATION") is not None)
            else []
        ),
        (
            expand(
                "{output_dir}/plots/classification/overlay/{segmentation_key}/{sample}.{segmentation_key}.{table_suffix}.gmm-overlay.png",
                output_dir=config["GENERAL"]["output_dir"],
                sample=SAMPLES,
                segmentation_key=SEGMENTATION_KEYS,
                table_suffix=config.get("SEGMENTATION", {}).get("tables_suffix"),
            )
            if (config.get("CLASSIFICATION") is not None)
            and (config.get("SEGMENTATION") is not None)
            else []
        ),
        # Tissue segmentation
        expand(
            "{output_dir}/done/tissue_segmentation/{sample}.{tissue_segmentation_key}",
            output_dir=config["GENERAL"]["output_dir"],
            sample=SAMPLES,
            tissue_segmentation_key=TISSUE_SEGMENTATION_KEYS,
        ),
        # HE patch embedding
        expand(
            "{output_dir}/done/patch_embedding/{sample}.{patch_embedding_key}",
            output_dir=config["GENERAL"]["output_dir"],
            sample=SAMPLES,
            patch_embedding_key=HE_PATCH_EMBEDDING_KEYS,
        ),
        # HE cell segmentation
        expand(
            "{output_dir}/done/he_cell_segmentation/{sample}.{he_cell_segmentation_key}",
            output_dir=config["GENERAL"]["output_dir"],
            sample=SAMPLES,
            he_cell_segmentation_key=HE_CELL_SEGMENTATION_KEYS,
        ),
        # HE tissue qc
        expand(
            "{output_dir}/done/he_tissue_qc/{sample}.{he_tissue_qc_key}",
            output_dir=config["GENERAL"]["output_dir"],
            sample=SAMPLES,
            he_tissue_qc_key=HE_TISSUE_QC_KEYS,
        ),
    default_target: True
