import os
import re

import pandas as pd


# Settings: the pipeline's defaults, then the run's config/config.yml when there is one, then
# any --configfile / --config given to Snakemake (`ophist run CONFIG` passes CONFIG this way).
# A run's config only needs the settings it changes.

configfile: f"{workflow.basedir}/config/config.yml"
if os.path.exists("config/config.yml"):
    configfile: "config/config.yml"


# Setup sample list (or single sample name)

samples_cfg = config["samples"]

if isinstance(samples_cfg, str) and not samples_cfg.endswith(".tsv"):
    SAMPLES = [samples_cfg]
    SAMPLE_SHEET = None
else:
    # sample_name, path (Xenium outs/ directory), optional x/y_min/max crop and contrast_N columns
    SAMPLE_SHEET = pd.read_csv(samples_cfg, sep="\t", dtype={"sample_name": str, "path": str}).set_index("sample_name")
    SAMPLES = SAMPLE_SHEET.index.tolist()

wildcard_constraints:
    sample="|".join(re.escape(s) for s in SAMPLES)


# spatialRNA's output store: written by rule 03, or already there for prerun data.
# Everything downstream of spatialRNA reads its transcripts from here.

PRERUN = config["preprocessing"]["type"] == "prerun"
EMBEDDINGS = config["spatialrna"]["embeddings"]
EMBEDDING_ROWS = f"{EMBEDDINGS}/{{sample}}/x"   # a sample's row count, for mem_by_rows()

# Apptainer images (--sdm apptainer), one per envs/*.yml; each rule's container: sits
# next to its conda: env, which --sdm conda uses instead.
CONTAINERS = config["containers"]

# Smoothing and the Potts model (rules 06, 07) come from the duetist package installed in
# the spatial env, or from a duetist checkout when duetist.src is set (docs/config.md).
DUETIST_SRC = str(config.get("duetist", {}).get("src") or "")

# The report's heading and file: report.title, and output/<title>.html with anything unsafe
# in a file name replaced. No title gives the default heading in output/report.html.
REPORT_TITLE = str(config.get("report", {}).get("title") or "").removesuffix(".html")
REPORT_FILE = re.sub(r"[^\w.-]+", "_", REPORT_TITLE).strip("_") or "report"
REPORT_HTML = f"output/{REPORT_FILE}.html"


# Condition DGE targets: only with several samples and a contrast_N column in samples.tsv

def has_condition_contrasts():
    header = pd.read_csv(samples_cfg, sep="\t", nrows=0).columns
    return any(re.match(r"^contrast_\d+$", c) for c in header)

def get_conditional_targets():
    if len(SAMPLES) > 1 and has_condition_contrasts():
        return ["output/condition_dge/condition_de.csv", "output/condition_dge/condition_gsea.csv",
                "output/condition_dge/condition_de_full.csv", "output/condition_dge/condition_gsea_full.csv"]
    return []


# Define all targets

rule all:
    input:
        expand("output/{sample}/lumen_plot.png", sample=SAMPLES),
        expand("output/{sample}/tx_lumen_plot.png", sample=SAMPLES),
        "output/plots/lumen_detection.png",
        expand("output/{sample}/smoothed_embeddings.png", sample=SAMPLES),
        expand("output/{sample}/hmrf_labels.png", sample=SAMPLES),
        [f"output/plots/hmrf_{p}.png" for p in ("domains", "niches", "domain_profiles", "history", "composition")],
        "output/domain_dge/domain_de.csv",
        "output/domain_dge/domain_gsea.csv",
        REPORT_HTML,
        get_conditional_targets()

localrules: all


# Resource allocation for job submission

include: "rules/resources.smk"


# Rules proper

# 1. Preprocessing - including splitting of multiplexed samples, and filtering of low-quality/negative control spots
include: "rules/01_preprocessing.smk"
# 2. Tissue detection - Masking of the tissue area to exclude background spots
include: "rules/02_tissue_detection.smk"
# 3. SpatialRNA - production of per-transcript embeddings reflecting their local spatial context
if not PRERUN:
    include: "rules/03_spatialrna.smk"
# 4. Lumen detection - detection of lumen in the tissue
include: "rules/04_lumen_detection.smk"
# 6. Smoothing - smoothing of the spatial embeddings onto a grid
include: "rules/06_smoothing.smk"
# 7. Niches and domains - a hierarchical Potts model of niches and domains across the cohort
include: "rules/07_hmrf_potts.smk"
# 8. Domain analysis - pseudobulk counts per domain, domain marker genes and GSEA
include: "rules/08_domain_analysis.smk"
# 9. DGE - differential gene expression analysis between conditions, if multiple samples and contrasts are provided
include: "rules/09_dge.smk"
# 10. Report - generation of a report summarizing the results
include: "rules/10_report.smk"
