##################################################################
###                                                            ###
### This snakemake allows to compute a full pipeline analyze   ###
### with hicexplorer. At the end, it will return HiC maps      ###
### and TADs plot for each resolutions, protocols and samples  ###
### given in config.yaml.                                      ###
###                                                            ###
### This snakemake is called to submit on cluster in           ###
### hicexplorer.sh but you can run it with the following cmd : ###
###                                                            ###
### snakemake -j 1 -p                                          ###
###                                                            ###
##################################################################

import os
import glob
import re
import pandas as pd

# Function definition
def set_run_names(sampled_protocols):
    ''' For each protocol retrieve the associated run names '''
    runs = []
    for protocol in sampled_protocols.index:
        matrix_pref = sampled_protocols.loc[protocol]['matrixprefix']
        prot_runs = []
        for f in glob.glob(matrix_pref+'*.matrix'):
            prot_runs.append(re.sub('_\d+.matrix', '', os.path.basename(f),))
        runs.append(",".join(set(prot_runs)))
    sampled_protocols['runs'] = runs
    return sampled_protocols


def get_raw_matrix(wildcards):
    """ input function for rule conversion
        returns the hicpro raw matrix associated to a protocol-run-resolution
    """
    mdir = sampled_protocols.loc[(wildcards.sample, wildcards.protocol)]['matrixprefix'].split('trio')[0]
    matrix_file = wildcards.run + "_" + wildcards.resolution + ".matrix"
    return mdir + matrix_file


def get_bed_file(wildcards):
    """ input function for rule conversion
        returns the hicpro bed file associated to a protocol-run-resolution
    """
    mdir = sampled_protocols.loc[(wildcards.sample, wildcards.protocol)]['matrixprefix'].split('trio')[0]
    bed_file = wildcards.run + "_" + wildcards.resolution + "_abs.bed"
    return mdir + bed_file


def get_matrices(wildcards):
    """ input function for rule mergeprotocol
        returns list of the h5 matrices associated with a protocol-resolution
    """
    runs = sampled_protocols.loc[(wildcards.sample, wildcards.protocol)]['runs'].split(',')
    matrices = []
    for run in runs:
        m = run + "_" + wildcards.resolution + ".h5"
        m_path = os.path.join("hdf5", wildcards.sample, wildcards.protocol, m)
        matrices.append(m_path)
    return matrices

def get_samples_matrices(wildcards):
    """ input function for rule merged
        return list of the merged h5 matrices from rule mergeprotocol from a association individual-resolution
    """
    matrices = []
    for prot in sampled_protocols['protocol']:
        if prot in sampled_protocols.loc[wildcards.sample]['protocol']:
            m_path = os.path.join("hdf5", wildcards.sample, prot, wildcards.resolution, "matrix.h5")
            if m_path not in matrices:
                matrices.append(m_path)
    return matrices


configfile: "config.yaml"

sampled_protocols = pd.read_csv(config["sampled_protocols"], sep = "\t").set_index(["sample","protocol"], drop = False)

sampled_protocols = set_run_names(sampled_protocols)

chromosomes = config['chromosomes']
resolutions = config['resolutions']
regions = config['regions']

protocols_and_merged = list(set(sampled_protocols['protocol'])) + ['Merged']

samples = config['samples']

# workdir: config['workdir']

wildcard_constraints:
    protocol="|".join(protocols_and_merged)

rule all:
        input:
            expand("plots/{sample}/merged/{resolution}/{region}_plot.png", sample=samples, protocol=protocols_and_merged, resolution=resolutions, region=regions),
            expand("TADs/{sample}/{resolution}/Done", sample=samples, resolution=resolutions),
            expand("TADs/plots/TADs_{sample}_{resolution}_{region}.png", sample=samples, resolution=resolutions, region=regions)

rule conversion:
    input:
        raw = get_raw_matrix,
        bed = get_bed_file
    output:
        "hdf5/{sample}/{protocol}/{run}_{resolution}.h5"
    shell:
        "hicConvertFormat --matrices {input.raw} --bedFileHicpro {input.bed}"
        " --outFileName {output} --outputFormat h5 --inputFormat hicpro"

rule mergeprotocol:
    input:
        get_matrices
    output:
        "hdf5/{sample}/{protocol}/{resolution}/matrix.h5"
    shell:
        "hicSumMatrices --matrices  {input} --outFileName {output}"

rule merge:
    input:
        # expand("hdf5/{{sample}}/{unit.protocol}/{{resolution}}/matrix.h5", unit=sampled_protocols.itertuples())
        get_samples_matrices
    output:
        "hdf5/{sample}/merged/{resolution}/matrix.h5"
    shell:
        "hicSumMatrices --matrices {input} --outFileName {output}"

rule adjust:
    input:
        "hdf5/{sample}/merged/{resolution}/matrix.h5",
    output:
        "hdf5/{sample}/merged/{resolution}/adjusted.h5"
    params:
        chromosomes = chromosomes
    shell:
        "hicAdjustMatrix --matrix {input} --outFileName {output} "
        " --chromosomes {params.chromosomes} --action keep ;"

rule normalize:
    input:
        "hdf5/{sample}/merged/{resolution}/adjusted.h5"
    output:
        "hdf5/{sample}/merged/{resolution}/normalized.h5"
    shell:
        "hicNormalize --matrices {input} --outFileName {output}"
        " --normalize smallest"

rule correct:
    input:
        "hdf5/{sample}/merged/{resolution}/normalized.h5"
    output:
        "hdf5/{sample}/merged/{resolution}/corrected.h5"
    shell:
        "hicCorrectMatrix correct --matrix {input} --outFileName {output}"

rule diagnostic:
    input:
        "hdf5/{sample}/merged/{resolution}/{type}.h5"
    output:
        "hdf5/{sample}/merged/{resolution}/{type}_diagnostic.png"
    shell:
        "hicCorrectMatrix diagnostic_plot -m {input} -o {output}"

rule plotMatrix:
    input:
        "hdf5/{sample}/merged/{resolution}/corrected.h5"
    output:
        "plots/{sample}/merged/{resolution}/{region}_plot.png"
    shell:
        "hicPlotMatrix --matrix {input} --outFileName {output} --region 1:{wildcards.region} --log"

rule findTADs:
    input:
        "hdf5/{sample}/merged/{resolution}/corrected.h5"
    output:
        "TADs/{sample}/{resolution}/Done"
    params:
        chr1_prefix = "TADs/{sample}/{resolution}/TADs_chr1",
        genomic_prefix = "TADs/{sample}/{resolution}/TADs_genomes",
    shell:
        "hicFindTADs --matrix {input} --outPrefix {params.chr1_prefix} --correctForMultipleTesting fdr --chromosomes 1 ; \
         hicFindTADs --matrix {input} --outPrefix {params.genomic_prefix} --correctForMultipleTesting fdr --chromosomes {chromosomes} ; \
         touch {output}"
         
rule tracksINI:
    input:
        matrix = expand("hdf5/{{sample}}/merged/{{resolution}}/corrected.h5"),
        domains = expand("TADs/{{sample}}/{{resolution}}/TADs_chr1_domains.bed"),
        score = expand("TADs/{{sample}}/{{resolution}}/TADs_chr1_score.bedgraph")
    params:
        title = "{sample}_{resolution}"
    output:
        "TADs/tracks/{sample}_{resolution}.ini"
    shell: 
        "printf '[x-axis] \n \
fontsize=10 \n \
\n \
[hic matrix] \n \
file = {input.matrix} \n \
title = {params.title} \n \
# depth is the maximum distance plotted in bp. In Hi-C tracks \n \
# the height of the track is calculated based on the depth such \n \
# that the matrix does not look deformed \n \
colormap = Spectral_r \n \
depth = 1000000 \n \
transform = log1p \n \
file_type = hic_matrix \n \
\n \
[tads] \n \
file = {input.domains} \n \
file_type = domains \n \
border_color = black \n \
color = none \n \
# the tads are overlay over the hic-matrix \n \
# the share-y options sets the y-axis to b e shared \n \
# between the Hi-C matrix and the TADs. \n \
overlay_previous = share-y \n \
\n \
[bedgraph] \n \
file = {input.score} \n \
color = blue \n \
height = 3 \n \
title = TAD score \n \
max_value = 10 \n \
\n \
[spacer] \n \
height = 1' > {output}"
        
rule plotTADs:
    input:
        expand("TADs/tracks/{{sample}}_{{resolution}}.ini")
    output:
        "TADs/plots/TADs_{sample}_{resolution}_{region}.png"
    shell:
        "hicPlotTADs --tracks {input} -o {output} --region 1:{wildcards.region}"

