# ------------------------------------------------------------------------------
# Python specific integration rule.
# TODO: decorate with some detailed information
# ------------------------------------------------------------------------------
rule integration_run_python:
    input:
        adata=expand(rules.integration_prepare.output, prep="h5ad", allow_missing=True)
    output:
        cfg.ROOT / "{scenario}/integration/{scaling}/{hvg}/{method}.h5ad"
    message:
        """
        Run {wildcards.method} on {wildcards.scaling} data
        feature selection: {wildcards.hvg}
        dataset: {wildcards.scenario}
        command: {params.cmd}
        hvgs: {params.hvgs}
        cell type option: {params.cell_type}
        output: {output[0]}
        """
    params:
        batch_key = lambda wildcards: cfg.get_from_scenario(wildcards.scenario, key="batch_key"),
        cell_type = cfg.get_celltype_option_for_integration,
        hvgs      = cfg.get_hvg,
        timing    = "-t" if cfg.timing else "",
        cmd       = f"conda run -n {cfg.py_env} python"
    benchmark:
        str(cfg.ROOT / "{scenario}/integration/{scaling}/{hvg}/{method}.h5ad.benchmark")
    shell:
        """
        {params.cmd} scripts/integration/runIntegration.py \
          -i {input.adata} -o {output} \
	      -b {params.batch_key} --method {wildcards.method} {params.hvgs} {params.cell_type} \
	      {params.timing}
        """

# ------------------------------------------------------------------------------
# R specific integration rule.
# TODO: decorate with some detailed information
# ------------------------------------------------------------------------------
rule integration_run_r:
    input:
        adata=expand(rules.integration_prepare.output, prep="RDS", allow_missing=True)
    output: cfg.ROOT / "{scenario}/integration/{scaling}/{hvg}/R/{method}.RDS"
    message:
        """
        output: {output[0]}
        Run {wildcards.method} on {wildcards.scaling} data
        feature selection: {wildcards.hvg}
        dataset: {wildcards.scenario}
        command: {params.cmd}
        hvgs: {params.hvgs}
        """
    params:
        batch_key = lambda w: cfg.get_from_scenario(w.scenario, key="batch_key"),
        hvgs      = lambda w: cfg.get_hvg(w, rules.integration_prepare.output, prep="RDS"),
        cmd       = f"conda run -n {cfg.r_env} Rscript",
        timing    = "-t" if cfg.timing else ""
    benchmark:
        str(cfg.ROOT / "{scenario}/integration/{scaling}/{hvg}/R/{method}.RDS.benchmark")
    shell:
        """
        {params.cmd} scripts/integration/runMethods.R -i {input.adata} \
            -o {output} -b {params.batch_key} \
            --method {wildcards.method} {params.hvgs} {params.timing}
        """


# ------------------------------------------------------------------------------
# Simply converts the RDS files created by the R scripts to h5ad files for
# further processing with the metrics rule
# ------------------------------------------------------------------------------
rule convert_RDS_h5ad:
    input:
        i = rules.integration_run_r.output
    output:
        cfg.ROOT / "{scenario}/integration/{scaling}/{hvg}/R/{method}.h5ad"
    message:
        """
        Convert integrated data from {wildcards.method} into h5ad
        """
    params:
        cmd = f"conda run -n {cfg.py_env} python"
    shell:
        """
        if [ {wildcards.method} == "conos" ]
        then
            {params.cmd} scripts/integration/runPost.py -i {input.i} -o {output} -c
        else
            {params.cmd} scripts/integration/runPost.py -i {input.i} -o {output}
        fi
        """


int_func, py_wildcards = cfg.get_all_wildcards(methods=cfg.get_all_python_methods())
rule integrated_python:
    input: expand(rules.integration_run_python.output, int_func, **py_wildcards)
    message: "Python integration methods"


int_func, r_wildcards = cfg.get_all_wildcards(methods=cfg.get_all_R_methods())
rule integrated_R:
    input: expand(rules.convert_RDS_h5ad.output, int_func, **r_wildcards)
    message: "R integration methods"

rule integration:
    input:
        rules.integrated_python.input,
        rules.integrated_R.input
    message: "Integration done"

51c193be618997542f5c85712f1bcb00f3244054
# ------------------------------------------------------------------------------
# Python specific integration rule.
# TODO: decorate with some detailed information
# ------------------------------------------------------------------------------
rule integration_run_python:
    input:
        adata=expand(rules.integration_prepare.output, prep="h5ad", allow_missing=True)
    output:
        cfg.ROOT / "{scenario}/integration/{scaling}/{hvg}/{method}.h5ad"
    message:
        """
        Run {wildcards.method} on {wildcards.scaling} data
        feature selection: {wildcards.hvg}
        dataset: {wildcards.scenario}
        command: {params.cmd}
        hvgs: {params.hvgs}
        cell type option: {params.cell_type}
        output: {output[0]}
        """
    params:
        batch_key = lambda wildcards: cfg.get_from_scenario(wildcards.scenario, key="batch_key"),
        cell_type = cfg.get_celltype_option_for_integration,
        hvgs      = cfg.get_hvg,
        timing    = "-t" if cfg.timing else "",
        cmd       = f"conda run -n {cfg.py_env} python"
    benchmark:
        str(cfg.ROOT / "{scenario}/integration/{scaling}/{hvg}/{method}.h5ad.benchmark")
    shell:
        """
        {params.cmd} scripts/integration/runIntegration.py \
          -i {input.adata} -o {output} \
	      -b {params.batch_key} --method {wildcards.method} {params.hvgs} {params.cell_type} \
	      {params.timing}
        """

# ------------------------------------------------------------------------------
# R specific integration rule.
# TODO: decorate with some detailed information
# ------------------------------------------------------------------------------
rule integration_run_r:
    input:
        adata=expand(rules.integration_prepare.output, prep="RDS", allow_missing=True)
    output: cfg.ROOT / "{scenario}/integration/{scaling}/{hvg}/R/{method}.RDS"
    message:
        """
        output: {output[0]}
        Run {wildcards.method} on {wildcards.scaling} data
        feature selection: {wildcards.hvg}
        dataset: {wildcards.scenario}
        command: {params.cmd}
        hvgs: {params.hvgs}
        """
    params:
        batch_key = lambda w: cfg.get_from_scenario(w.scenario, key="batch_key"),
        hvgs      = lambda w: cfg.get_hvg(w, rules.integration_prepare.output, prep="RDS"),
        cmd       = f"conda run -n {cfg.r_env} Rscript",
        timing    = "-t" if cfg.timing else ""
    benchmark:
        str(cfg.ROOT / "{scenario}/integration/{scaling}/{hvg}/R/{method}.RDS.benchmark")
    shell:
        """
        {params.cmd} scripts/integration/runMethods.R -i {input.adata} \
            -o {output} -b {params.batch_key} \
            --method {wildcards.method} {params.hvgs} {params.timing}
        """


# ------------------------------------------------------------------------------
# Simply converts the RDS files created by the R scripts to h5ad files for
# further processing with the metrics rule
# ------------------------------------------------------------------------------
rule convert_RDS_h5ad:
    input:
        i = rules.integration_run_r.output
    output:
        cfg.ROOT / "{scenario}/integration/{scaling}/{hvg}/R/{method}.h5ad"
    message:
        """
        Convert integrated data from {wildcards.method} into h5ad
        """
    params:
        cmd = f"conda run -n {cfg.py_env} python"
    shell:
        """
        if [ {wildcards.method} == "conos" ]
        then
            {params.cmd} scripts/integration/runPost.py -i {input.i} -o {output} -c
        else
            {params.cmd} scripts/integration/runPost.py -i {input.i} -o {output}
        fi
        """


int_func, py_wildcards = cfg.get_all_wildcards(methods=cfg.get_all_python_methods())
rule integrated_python:
    input: expand(rules.integration_run_python.output, int_func, **py_wildcards)
    message: "Python integration methods"


int_func, r_wildcards = cfg.get_all_wildcards(methods=cfg.get_all_R_methods())
rule integrated_R:
    input: expand(rules.convert_RDS_h5ad.output, int_func, **r_wildcards)
    message: "R integration methods"

rule integration:
    input:
        rules.integrated_python.input,
        rules.integrated_R.input
    message: "Integration done"
