#### Snakefile for the up-stream analysis of snRNA-seq data
#### Developer: Pengwei Dong, Shitong Ding, Guanlin Wang

#### We recommand you create a new conda environment with Cell Ranger and CellBender installed
## Arguments file ---- you can set arguments here

configfile: "config.yaml"

## Sample names
SAMPLES = config["samples"]

#### Final aim ---- create filtered.h5 files for Seurat/Scanpy analysis
rule all:
    input:
        expand("output/cellbender/{sample}_filtered.h5", sample=SAMPLES)

## Step 1: Cell Ranger count
rule cellranger_count:
    input:
        fastq_dir="data/{sample}"  # Directory where FastQ files are located
    output:
        matrix_dir="output/cellranger/{sample}/outs/raw_feature_bc_matrix.h5"  # Raw h5 files in cellranger outputs
    log:
        "logs/{sample}_cellranger_count.log"
    params:
        transcriptome=config["cellranger"]["transcriptome"]
    shell:
        """
        ./cellranger-7.1.0/cellranger count --id={wildcards.sample} \
                                            --transcriptome={params.transcriptome} \
                                            --fastqs={input.fastq_dir} \
                                            --sample={wildcards.sample} > {log} 2>&1
        cp {wildcards.sample}/outs/raw_feature_bc_matrix.h5 {output.matrix_dir}
        """

## Step 2: Removing ambient RNAs by applying CellBender
rule cellbender_remove_background:
    input:
        matrix="output/cellranger/{sample}/outs/raw_feature_bc_matrix.h5"
    output:
        h5="output/cellbender/{sample}_filtered.h5"
    params:
        total_droplet_cells=config["cellbender"]["total_droplet_cells"],
        expected_cells=config["cellbender"]["expected_cells"],
        epochs=config["cellbender"]["epochs"]
    log:
        "logs/{sample}_cellbender.log"
    shell:
        """
        cellbender remove-background \
                   --input {input.matrix} \
                   --output {output.h5} \
                   --total-droplets-included {params.total_droplet_cells} \
                   --expected-cells {params.expected_cells} \
                   --epochs {params.epochs} 
                   --cuda > {log} 2>&1  ### --cuda flag solely for the purposes of being able to run this command on a CPU
        """
