Snakefile

###### Import libraries #######
import os
from os import path
from glob import glob
import itertools
import cooler
import numpy as np
import pandas as pd
import bioframe
import toolz
from cooltools.lib.io import read_viewframe_from_file
from scipy import spatial


localrules:
    all,
    make_pentads,
    make_diff_pentads,
    merge_dots_across_resolutions,
    make_differential_insulation,
    make_tads,
    call_dots_mustache,
    get_bed_data,
    get_cool_data,


###### Read config parameters #######

project_folder = config.get("project_folder", "results")
inputs_folder = config.get("inputs_folder", "inputs")
coolers_folder = config.get("coolers_folder", path.join(inputs_folder, "coolers"))
beds_folder = path.normpath(config.get("beds_folder", path.join(inputs_folder, "beds")))
beds_folder_name = path.basename(beds_folder)
bedpes_folder = path.normpath(
    config.get("bedpes_folder", path.join(inputs_folder, "bedpes"))
)
bedpes_folder_name = path.basename(bedpes_folder)
expected_folder = path.normpath(
    config.get("expected_folder", path.join(project_folder, "expected"))
)
pileups_folder = path.normpath(
    config.get("pileups_folder", path.join(project_folder, "pileups"))
)
eigenvectors_folder = path.normpath(
    config.get("eigenvectors_folder", path.join(project_folder, "eigenvectors"))
)
saddles_folder = path.normpath(
    config.get("saddles_folder", path.join(project_folder, "saddles"))
)
insulation_folder = path.normpath(
    config.get("insulation_folder", path.join(project_folder, "insulation"))
)
tad_folder = path.normpath(config.get("tad_folder", path.join(project_folder, "tads")))
tad_folder_name = path.basename(tad_folder)
loop_folder = path.normpath(config.get("dot_folder", path.join(project_folder, "dots")))
loop_folder_name = path.basename(loop_folder)
boundary_folder = path.normpath(
    config.get("boundary_folder", path.join(project_folder, "boundaries"))
)
boundary_folder_name = path.basename(boundary_folder)


include: "rules/common.smk"


### Input genome
genome = config["genome"]
chroms = pd.read_csv(
    f'{config["chromsizes"]}', sep="\t", names=["chrom", "start", "end"]
)["chrom"].values

### Input cool files
from urllib.parse import urlparse


def will_download(link):
    parsed_path = urlparse(link)
    return False if parsed_path.scheme == "" else True


try:
    samples_df = pd.read_csv(
        config["samples"],
        sep="\t",
        header=0,
        comment="#",
        dtype={"do_dots": bool, "do_tads": bool},
    )
except:
    raise ValueError(
        "Could not read file with samples, please ensure it exists and has data in it"
    )
if "sample" not in samples_df.columns:
    raise ValueError(
        'Column "sample" has to be in the file with description of samples'
    )
if "file" not in samples_df.columns:
    raise ValueError('Column "file" has to be in the file with description of samples')

samples_df = samples_df.fillna(value=False)
samples_df.loc[:, "will_download"] = samples_df.file.apply(will_download)
samples_df.loc[:, "local_path"] = samples_df.apply(
    lambda x: f"{coolers_folder}/{x['sample']}.mcool" if x.will_download else x["file"],
    axis=1,
)
samples_df = samples_df.set_index("sample")

samples = list(samples_df.index)
coollinks_dict = dict(
    samples_df.query("will_download")["file"]
)  # dict with cools to be downloaded
coolfiles_dict = dict(samples_df["local_path"])

for coolfile in coolfiles_dict.keys():
    if coolfile not in coollinks_dict.keys():
        verify_view_cooler(
            cooler.Cooler(f"{coolfile}::resolutions/{config['resolutions'][0]}")
        )

# Setting up pairwise comparisons
sample_pairs = []

if config["fields_to_differ"] is not None:
    for name, group in samples_df.groupby(
        lambda x: config["fields_to_match"]
        if config["fields_to_match"] is not None
        else True
    ):
        for sample1, sample2 in itertools.permutations(group.index, 2):
            if np.all(
                samples_df.loc[
                    sample1, list(config["fields_to_differ"].keys())
                ].to_numpy()
                != np.array(list(config["fields_to_differ"].values()))
            ) and np.all(
                samples_df.loc[
                    sample2, list(config["fields_to_differ"].keys())
                ].to_numpy()
                == np.array(list(config["fields_to_differ"].values()))
            ):
                sample_pairs.append((sample1, sample2))
sample_pairs_string = [f"{s1}_vs_{s2}" for s1, s2 in sample_pairs]
sample_pairs_string_with_reverse = sample_pairs_string + [
    f"{s2}_vs_{s1}" for s1, s2 in sample_pairs
]

# Setting up annotation files
bedfiles_local = get_files(beds_folder, "bed")
bedpefiles_local = get_files(bedpes_folder, "bedpe")

local_bed_names = {
    path.splitext(bedfile)[0]: f"{beds_folder}/{bedfile}" for bedfile in bedfiles_local
}
local_bedpe_names = {
    path.splitext(bedpefile)[0]: f"{bedpes_folder}/{bedpefile}"
    for bedpefile in bedpefiles_local
}

if config["annotations"]:
    try:
        bed_df = pd.read_csv(config["annotations"], sep="\t", header=0, comment="#")
    except:
        raise ValueError(
            "Could not read file with samples, please ensure it exists and has data in it"
        )
    bed_df.loc[:, "will_download"] = bed_df.file.apply(will_download)
    bed_df.loc[:, "local_path"] = bed_df.apply(
        lambda x: make_local_path(x.bedname, x.format) if x.will_download else x.file,
        axis=1,
        result_type="reduce",
    )
    bed_df = bed_df.set_index("bedname").replace("-", np.nan)
else:
    bed_df = pd.DataFrame(
        columns=["bedname", "file", "format", "will_download", "local_path"]
        + list(config["pileups"]["arguments"])
    )

pileup_params = config["pileups"]["arguments"]

bed_df[list(pileup_params.keys())] = ~bed_df[list(pileup_params.keys())].isna()

bedlinks_dict = dict(
    bed_df.loc[bed_df["will_download"], "file"]
)  # dict with beds to be downloaded
bedfiles_dict = dict(bed_df["local_path"])
bedfiles_dict.update(local_bed_names)
bedfiles_dict.update(local_bedpe_names)
bedfiles = list(bedfiles_dict.keys())
# bedfiles_pileups = [bf for bf in bedfiles if bed_df.loc[bf, 'pileups']]
bedtype_dict = dict(bed_df["format"])
# bedpe_pileups_mindist, bedpe_pileups_maxdist = config['bedpe_pileups_distance_limits']

if config["samples_annotations_combinations"]:
    try:
        samples_annotations = ~pd.read_csv(
            config["samples_annotations_combinations"],
            sep="\t",
            header=0,
            index_col=0,
            comment="#",
        ).isna()
    except:
        raise ValueError(
            "Could not read file with sample-annotation combinations,"
            "please ensure it exists and has data in it"
        )
else:
    samples_annotations = pd.DataFrame(
        np.ones((len(samples), len(bedfiles))), index=samples, columns=bedfiles
    ).astype(bool)

### Data resolutions
if config["eigenvector"]["do"]:
    eigenvector_resolutions = config["eigenvector"]["resolutions"]

if config["saddle"]["do"]:
    saddle_mindist, saddle_maxdist = config["saddle"]["distance_limits"]
    saddle_mindists = [
        int(saddle_mindist * 2**i)
        for i in np.arange(0, np.log2(saddle_maxdist / saddle_mindist))
    ]
    saddle_separations = [f"_dist_{mindist}-{mindist*2}" for mindist in saddle_mindists]

if config["pileups"]["do"]:
    shifts = config["pileups"]["shifts"]
    pileup_norms = []
    if shifts > 0:
        pileup_norms.append(f"{shifts}-shifts")
    if config["pileups"]["expected"]:
        pileup_norms.append("expected")
    if len(pileup_norms) == 0:
        raise ValueError("Please use expected or shifts to normalize pileups")
    pileups_mindist, pileups_maxdist = config["pileups"]["distance_limits"]
    pileup_resolutions = config["pileups"]["resolutions"]
    mindists = [
        int(pileups_mindist * 2**i)
        for i in np.arange(0, np.log2(pileups_maxdist / pileups_mindist))
    ]
    separations = [f"_dist_{mindist}-{mindist*2}" for mindist in mindists]


if config["insulation"]["do"] and config["call_TADs"]["do"]:
    config["insulation"]["resolutions"] = merge_dicts(
        config["insulation"]["resolutions"], config["call_TADs"]["resolutions"]
    )
elif config["call_TADs"]["do"]:
    config["insulation"]["resolutions"] = config["call_TADs"]["resolutions"]

if config["insulation"]["do"]:
    insul_res_win_comb = []
    for resolution in config["insulation"]["resolutions"]:
        for win in config["insulation"]["resolutions"][resolution]:
            insul_res_win_comb.append(f"{resolution}_{win}")

if config["call_TADs"]["do"]:
    tad_res_win_comb = []
    for resolution in config["call_TADs"]["resolutions"]:
        for win in config["call_TADs"]["resolutions"][resolution]:
            tad_res_win_comb.append(f"{resolution}_{win}")

# chroms = cooler.Cooler(f'{coolers_folder}/{coolfiles[0]}::resolutions/{resolutions[0]}').chroms()[:]

# bedpe_mindists = [int(bedpe_pileups_mindist*2**i) for i in np.arange(0, np.log2(bedpe_pileups_maxdist/bedpe_pileups_mindist))]
# bedpe_separations = [f'{mindist}-{mindist*2}' for mindist in bedpe_mindists]

###### Setup outputs for the pipeline
# Expected files
expecteds = []
if config["expected"]["do"]:
    if config["expected"]["cis"]:
        expecteds += expand(
            f"{expected_folder}/{{sample}}_{{resolution}}.expected.tsv",
            sample=samples,
            resolution=config["expected"]["resolutions"],
        )
    if config["expected"]["trans"]:
        expecteds += expand(
            f"{expected_folder}/{{sample}}_{{resolution}}.expected.trans.tsv",
            sample=samples,
            resolution=config["expected"]["resolutions"],
        )

# Insulation, boundaries, TADs and their pileups
insulation = (
    expand(
        f"{insulation_folder}/{{sample}}_{{resolution}}.insulation.tsv",
        sample=samples,
        resolution=config["insulation"]["resolutions"].keys(),
    )
    if (config["insulation"]["do"] or config["call_TADs"]["do"])
    else []
)

boundaries = (
    expand(
        f"{boundary_folder}/Boundaries_{{sample}}_{{resolution_window}}.bed",
        sample=samples,
        resolution_window=insul_res_win_comb,
    )
    if (config["insulation"]["do"] or config["call_TADs"]["do"])
    else []
)

boundaries_pileups = (
    expand(
        f"{pileups_folder}/{boundary_folder_name}/{{sample}}-{{resolution}}_over_Boundaries_{{sample}}_{{resolution_window}}_{{norm}}_local.clpy",
        sample=samples,
        resolution=pileup_resolutions,
        resolution_window=insul_res_win_comb,
        norm=pileup_norms,
    )
    if (config["insulation"]["do"] and config["insulation"]["pileup"])
    else []
)

diff_boundaries = (
    expand(
        f"{boundary_folder}/Diff_boundaries_{{comparison}}_{{resolution_window}}.bed",
        comparison=sample_pairs_string_with_reverse,
        resolution_window=insul_res_win_comb,
    )
    if config["compare_boundaries"]["do"]
    else []
)
diff_boundaries_pileups = (
    expand(
        f"{pileups_folder}/{boundary_folder_name}/{{sample}}-{{resolution}}_over_Diff_boundaries_{{comparison}}_{{resolution_window}}_{{norm}}_local.clpy",
        sample=samples,
        resolution=pileup_resolutions,
        comparison=sample_pairs_string_with_reverse,
        resolution_window=insul_res_win_comb,
        norm=pileup_norms,
    )
    if config["compare_boundaries"]["do"] and config["compare_boundaries"]["pileup"]
    else []
)

tads = (
    expand(
        f"{tad_folder}/TADs_{{sample}}_{{tad_res_win}}.bed",
        sample=samples_df.index[samples_df["do_tads"]],
        tad_res_win=tad_res_win_comb,
    )
    if config["call_TADs"]["do"]
    else []
)


tads_pileups = (
    expand(
        f"{pileups_folder}/{tad_folder_name}/{{sample}}-{{resolution}}_over_TADs_{{sampleTADs}}_{{tad_res_win}}_{{norm}}_local_rescaled.clpy",
        sample=samples,
        resolution=pileup_resolutions,
        sampleTADs=samples_df.index[samples_df["do_tads"]],
        tad_res_win=tad_res_win_comb,
        norm=pileup_norms,
    )
    if config["call_TADs"]["do"] and config["call_TADs"]["pileup"]
    else []
)

# Dots and dot pileups
dot_methods = [
    m for m in config["call_dots"]["methods"] if config["call_dots"]["methods"][m]["do"]
]
if dot_methods:
    dots = expand(
        f"{loop_folder}/merged_resolutions/Dots_{{method}}_{{sampleDots}}.bedpe",
        method=dot_methods,
        sampleDots=samples_df.index[samples_df["do_dots"]],
    )

    dots_pileups = (
        expand(
            f"{pileups_folder}/{loop_folder_name}/{{sample}}-{{resolution}}_over_Dots_{{method}}_{{sampleDots}}_{{norm}}_{{mode}}.clpy",
            sample=samples,
            resolution=pileup_resolutions,
            method=dot_methods,
            sampleDots=samples_df.index[samples_df["do_dots"]],
            norm=pileup_norms,
            mode=["distal", "by_distance"],
        )
        if config["call_dots"]["pileup"]
        else []
    )
else:
    dots = []
    dots_pileups = []

# Setup file types for dot bedpes files and boundary/TAD beds
for file in dots:
    name = path.splitext(path.basename(file))[0]
    bedfiles_dict[name] = file
    bedtype_dict[name] = "bedpe"
for file in tads + diff_boundaries:
    name = path.splitext(path.basename(file))[0]
    bedfiles_dict[name] = file
    bedtype_dict[name] = "bed"

# Pileups for annotation files
beds_pileups = []
if config["pileups"]["do"]:
    for bedname, row in bed_df.iterrows():
        modes = []
        for mode in pileup_params.keys():
            if row[mode]:
                modes += [mode]

        for sample in samples:
            if sample not in samples_annotations.index:
                continue
            if (
                bedname in samples_annotations.columns
                and not samples_annotations.loc[sample, bedname]
            ):
                continue
            beds_pileups += expand(
                f"{pileups_folder}/{beds_folder_name}/{sample}-{{resolution}}_over_{bedname}_{{norm}}_{{mode}}.clpy",
                resolution=pileup_resolutions,
                norm=pileup_norms,
                mode=modes,
            )

# Saddles and compartment beds
saddles = (
    expand(
        f"{saddles_folder}/{{sample}}_{{resolution}}_over_{{sample}}_eig_{{bins}}{{dist}}.{{ending}}",
        sample=samples,
        resolution=eigenvector_resolutions,
        bins=config["saddle"]["bins"],
        dist=saddle_separations + [""],
        ending=["saddledump.npz", "digitized.tsv"],
    )
    if config["saddle"]["do"]
    else []
)
for sample1, sample2 in sample_pairs:
    saddles += (
        expand(
            f"{saddles_folder}/{sample1}_{{resolution}}_over_{sample2}_eig_{{bins}}{{dist}}.{{ending}}",
            resolution=eigenvector_resolutions,
            bins=config["saddle"]["bins"],
            dist=saddle_separations + [""],
            ending=["saddledump.npz", "digitized.tsv"],
        )
        if config["saddle"]["do"]
        else []
    )

compartments = (
    expand(
        f"{eigenvectors_folder}/compartments/{{sample}}_{{resolution}}_compartments.{{mode}}.bed",
        sample=samples,
        resolution=eigenvector_resolutions,
        mode=["cis"],  # Add option to do trans in the future
    )
    if config["eigenvector"]["save_compartment_beds"]
    else []
)


# Pentads for all samples based on compartments from all samples - for now all combinations!
pentads = (
    expand(
        f"{eigenvectors_folder}/pentads/{{sample}}_{config['pentads']['data_resolution']}_over_compartments_{{sample}}_{config['pentads']['eigenvector_resolution']}_{{norm}}_pentads.clpy",
        sample=samples,
        norm=config["pentads"]["norms"],
        mode=config["pentads"]["modes"],  # Add option to do trans in the future
    )
    if config["pentads"]["do"]
    else []
)
for sample1, sample2 in sample_pairs:
    pentads += (
        expand(
            f"{eigenvectors_folder}/pentads/{sample1}_{config['pentads']['data_resolution']}_over_compartments_{sample2}_{config['pentads']['eigenvector_resolution']}_{{norm}}_pentads.clpy",
            norm=config["pentads"]["norms"],
            mode=config["pentads"]["modes"],  # Add option to do trans in the future
        )
        if config["pentads"]["do"] and config["pentads"]["do_diff"]
        else []
    )

# Differential pentads calculated using compartments from sample2
diff_pentads = []
if config["pentads"]["do_diff"]:
    for sample1, sample2 in sample_pairs:
        diff_pentads += expand(
            f"{eigenvectors_folder}/pentads/diff/{sample1}_vs_{sample2}_{config['pentads']['data_resolution']}_over_compartments_{sample2}_{config['pentads']['eigenvector_resolution']}_{{norm}}_pentads.clpy",
            norm=config["pentads"]["norms"],
        )


# Pentads for all samples based on compartments from all samples - for now all combinations!
pentads = (
    expand(
        f"{eigenvectors_folder}/pentads/{{sample}}_{config['pentads']['data_resolution']}_over_compartments_{{sample}}_{config['pentads']['eigenvector_resolution']}_{{norm}}_pentads.clpy",
        sample=samples,
        norm=config["pentads"]["norms"],
        mode=config["pentads"]["modes"],  # Add option to do trans in the future
    )
    if config["pentads"]["do"]
    else []
)
for sample1, sample2 in sample_pairs:
    pentads += (
        expand(
            f"{eigenvectors_folder}/pentads/{sample1}_{config['pentads']['data_resolution']}_over_compartments_{sample2}_{config['pentads']['eigenvector_resolution']}_{{norm}}_pentads.clpy",
            norm=config["pentads"]["norms"],
            mode=config["pentads"]["modes"],  # Add option to do trans in the future
        )
        if config["pentads"]["do"] and config["pentads"]["do_diff"]
        else []
    )

# Differential pentads calculated using compartments from sample2
diff_pentads = []
if config["pentads"]["do_diff"]:
    for sample1, sample2 in sample_pairs:
        diff_pentads += expand(
            f"{eigenvectors_folder}/pentads/diff/{sample1}_vs_{sample2}_{config['pentads']['data_resolution']}_over_compartments_{sample2}_{config['pentads']['eigenvector_resolution']}_{{norm}}_pentads.clpy",
            norm=config["pentads"]["norms"],
        )


###### Define rules #######
rule all:
    input:
        lambda wildcards: expecteds,
        lambda wildcards: diff_boundaries,
        lambda wildcards: diff_boundaries_pileups,
        lambda wildcards: tads,
        lambda wildcards: tads_pileups,
        lambda wildcards: dots,
        lambda wildcards: dots_pileups,
        lambda wildcards: beds_pileups,
        lambda wildcards: saddles,


include: "rules/pentads.smk"
include: "rules/dots.smk"
include: "rules/eigenvectors.smk"
include: "rules/insulation_tads.smk"
include: "rules/expected.smk"
include: "rules/saddle.smk"
include: "rules/pileups.smk"


rule get_cool_data:
    output:
        f"{coolers_folder}/{{sample}}.mcool",
    threads: 1
    resources:
        mem_mb=256,
        runtime=60,
    params:
        file=lambda wildcards: coollinks_dict[wildcards.sample],
    run:
        get_file(str(params.file), output[0])
        verify_view_cooler(
            cooler.Cooler(f"{output[0]}::resolutions/{config['resolutions'][0]}")
        )


rule get_bedpe_data:
    output:
        f"{bedpes_folder}/{{bedname}}.bedpe",
    threads: 1
    resources:
        mem_mb=256,
        runtime=60,
    params:
        file=lambda wildcards: bedlinks_dict[wildcards.bedname],
    run:
        get_file(str(params.file), str(output))


rule get_bed_data:
    output:
        f"{beds_folder}/{{bedname}}.bed",
    threads: 1
    resources:
        mem_mb=256,
        runtime=60,
    params:
        file=lambda wildcards: bedlinks_dict[wildcards.bedname],
    run:
        get_file(str(params.file), str(output))