<commit_before>"""Handle extraction of final files from processing pipelines into storage.
"""
import datetime
import os

import toolz as tz

from bcbio import log, utils
from bcbio.upload import shared, filesystem, galaxy, s3
from bcbio.pipeline import run_info
import bcbio.pipeline.datadict as dd

_approaches = {"filesystem": filesystem,
               "galaxy": galaxy,
               "s3": s3}

def from_sample(sample):
    """Upload results of processing from an analysis pipeline sample.
    """
    upload_config = sample.get("upload")
    if upload_config:
        approach = _approaches[upload_config.get("method", "filesystem")]
        for finfo in _get_files(sample):
            approach.update_file(finfo, sample, upload_config)
        for finfo in _get_files_project(sample, upload_config):
            approach.update_file(finfo, None, upload_config)
    return [[sample]]

# ## File information from sample

def _get_files(sample):
    """Retrieve files for the sample, dispatching by analysis type.

    Each file is a dictionary containing the path plus associated
    metadata about the file and pipeline versions.
    """
    analysis = sample.get("analysis")
    if analysis.lower() in ["variant", "snp calling", "variant2", "standard"]:
        return _get_files_variantcall(sample)
    elif analysis in ["RNA-seq"]:
        return _get_files_rnaseq(sample)
    elif analysis.lower() in ["chip-seq"]:
        return _get_files_chipseq(sample)
    elif analysis.lower() in ["sailfish"]:
        return _get_files_sailfish(sample)
    else:
        return []

def _get_files_sailfish(sample):
    out = []
    out.append({"path": sample["sailfish_dir"],
                "type": "directory",
                "ext": "sailfish"})
    return _add_meta(out, sample)

def _get_files_rnaseq(sample):
    out = []
    algorithm = sample["config"]["algorithm"]
    out = _maybe_add_summary(algorithm, sample, out)
    out = _maybe_add_alignment(algorithm, sample, out)
    out = _maybe_add_disambiguate(algorithm, sample, out)
    out = _maybe_add_counts(algorithm, sample, out)
    out = _maybe_add_cufflinks(algorithm, sample, out)
    out = _maybe_add_oncofuse(algorithm, sample, out)
    out = _maybe_add_rnaseq_variant_file(algorithm, sample, out)
    return _add_meta(out, sample)

def _get_files_chipseq(sample):
    out = []
    algorithm = sample["config"]["algorithm"]
    out = _maybe_add_summary(algorithm, sample, out)
    out = _maybe_add_alignment(algorithm, sample, out)
    return _add_meta(out, sample)

def _add_meta(xs, sample=None, config=None):
    out = []
    for x in xs:
        x["mtime"] = shared.get_file_timestamp(x["path"])
        if sample and "sample" not in x:
            if isinstance(sample["name"], (tuple, list)):
                name = sample["name"][-1]
            else:
                name = "%s-%s" % (sample["name"],
                                  run_info.clean_name(sample["description"]))
            x["sample"] = name
        if config:
            if "fc_name" in config and "fc_date" in config:
                x["run"] = "%s_%s" % (config["fc_date"], config["fc_name"])
            else:
                x["run"] = "project_%s" % datetime.datetime.now().strftime("%Y-%m-%d")
        out.append(x)
    return out

def _get_files_variantcall(sample):
    """Return output files for the variant calling pipeline.
    """
    out = []
    algorithm = sample["config"]["algorithm"]
    out = _maybe_add_summary(algorithm, sample, out)
    out = _maybe_add_alignment(algorithm, sample, out)
    out = _maybe_add_disambiguate(algorithm, sample, out)
    out = _maybe_add_variant_file(algorithm, sample, out)
    out = _maybe_add_sv(algorithm, sample, out)
    out = _maybe_add_validate(algorithm, sample, out)
    return _add_meta(out, sample)

def _maybe_add_validate(algorith, sample, out):
    for i, plot in enumerate(tz.get_in(("validate", "grading_plots"), sample, [])):
        ptype = os.path.splitext(plot)[-1].replace(".", "")
        out.append({"path": plot,
                    "type": ptype,
                    "ext": "validate%s" % ("" if i == 0 else "-%s" % (i + 1))})
    return out

def _maybe_add_rnaseq_variant_file(algorithm, sample, out):
    if sample.get("vrn_file"):
        out.append({"path": sample.get("vrn_file"),
                    "type": "vcf",
                    "ext": "vcf"})
    return out

def _maybe_add_variant_file(algorithm, sample, out):
    if sample.get("align_bam") is not None and sample.get("vrn_file"):
        for x in sample["variants"]:
            out.extend(_get_variant_file(x, ("vrn_file",)))
            if x.get("bed_file"):
                out.append({"path": x["bed_file"],
                            "type": "bed",
                            "ext": "%s-callregions" % x["variantcaller"],
                            "variantcaller": x["variantcaller"]})
            if x.get("vrn_stats"):
                for extra, fname in x["vrn_stats"].items():
                    ext = utils.splitext_plus(fname)[-1].replace(".", "")
                    out.append({"path": fname,
                                "type": ext,
                                "ext": "%s-%s" % (x["variantcaller"], extra),
                                "variantcaller": x["variantcaller"]})
    return out

def _maybe_add_sv(algorithm, sample, out):
    if sample.get("align_bam") is not None and sample.get("sv"):
        for svcall in sample["sv"]:
            out.extend(_get_variant_file(svcall, ("vrn_file",)))
            out.extend(_get_variant_file(svcall, ("bedpe_file",)))
            out.extend(_get_variant_file(svcall, ("cnr_bed",)))
            out.extend(_get_variant_file(svcall, ("cnr_bedgraph",)))
            if "sample_bed" in svcall:
                out.append({"path": svcall["sample_bed"],
                            "type": "bed",
                            "ext": "%s-sample" % svcall["variantcaller"],
                            "variantcaller": svcall["variantcaller"]})
            if "validate" in svcall:
                for vkey in ["csv", "plot", "df"]:
                    vfile = tz.get_in(["validate", vkey], svcall)
                    if vfile:
                        to_u = []
                        if isinstance(vfile, dict):
                            for svtype, fname in vfile.items():
                                to_u.append((fname, "-%s" % svtype))
                        else:
                            to_u.append((vfile, "-%s" % vkey if vkey in ["df"] else ""))
                        for vfile, ext in to_u:
                            vext = os.path.splitext(vfile)[-1].replace(".", "")
                            out.append({"path": vfile,
                                        "type": vext,
                                        "ext": "%s-sv-validate%s" % (svcall["variantcaller"], ext),
                                        "variantcaller": svcall["variantcaller"]})
            if "plot" in svcall:
                for plot_name, fname in svcall["plot"].items():
                    ext = os.path.splitext(fname)[-1].replace(".", "")
                    out.append({"path": fname,
                                "type": ext,
                                "ext": "%s-%s" % (svcall["variantcaller"], plot_name),
                                "variantcaller": svcall["variantcaller"]})
    return out

def _get_variant_file(x, key):
    """Retrieve VCF file with the given key if it exists, handling bgzipped.
    """
    out = []
    fname = utils.get_in(x, key)
    upload_key = list(key)
    upload_key[-1] = "do_upload"
    do_upload = tz.get_in(tuple(upload_key), x, True)
    if fname and do_upload:
        if fname.endswith(".vcf.gz"):
            out.append({"path": fname,
                        "type": "vcf.gz",
                        "ext": x["variantcaller"],
                        "variantcaller": x["variantcaller"]})
            if utils.file_exists(fname + ".tbi"):
                out.append({"path": fname + ".tbi",
                            "type": "vcf.gz.tbi",
                            "index": True,
                            "ext": x["variantcaller"],
                            "variantcaller": x["variantcaller"]})
        elif fname.endswith((".vcf", ".bed", ".bedpe", ".bedgraph")):
            ftype = utils.splitext_plus(fname)[-1][1:]
            out.append({"path": fname,
                        "type": ftype,
                        "ext": x["variantcaller"],
                        "variantcaller": x["variantcaller"]})
    return out

def _maybe_add_summary(algorithm, sample, out):
    out = []
    if "summary" in sample:<commit_msg>Include cnr and cns CNVkit files in final output directory<commit_after>"""Handle extraction of final files from processing pipelines into storage.
"""
import datetime
import os

import toolz as tz

from bcbio import log, utils
from bcbio.upload import shared, filesystem, galaxy, s3
from bcbio.pipeline import run_info
import bcbio.pipeline.datadict as dd

_approaches = {"filesystem": filesystem,
               "galaxy": galaxy,
               "s3": s3}

def from_sample(sample):
    """Upload results of processing from an analysis pipeline sample.
    """
    upload_config = sample.get("upload")
    if upload_config:
        approach = _approaches[upload_config.get("method", "filesystem")]
        for finfo in _get_files(sample):
            approach.update_file(finfo, sample, upload_config)
        for finfo in _get_files_project(sample, upload_config):
            approach.update_file(finfo, None, upload_config)
    return [[sample]]

# ## File information from sample

def _get_files(sample):
    """Retrieve files for the sample, dispatching by analysis type.

    Each file is a dictionary containing the path plus associated
    metadata about the file and pipeline versions.
    """
    analysis = sample.get("analysis")
    if analysis.lower() in ["variant", "snp calling", "variant2", "standard"]:
        return _get_files_variantcall(sample)
    elif analysis in ["RNA-seq"]:
        return _get_files_rnaseq(sample)
    elif analysis.lower() in ["chip-seq"]:
        return _get_files_chipseq(sample)
    elif analysis.lower() in ["sailfish"]:
        return _get_files_sailfish(sample)
    else:
        return []

def _get_files_sailfish(sample):
    out = []
    out.append({"path": sample["sailfish_dir"],
                "type": "directory",
                "ext": "sailfish"})
    return _add_meta(out, sample)

def _get_files_rnaseq(sample):
    out = []
    algorithm = sample["config"]["algorithm"]
    out = _maybe_add_summary(algorithm, sample, out)
    out = _maybe_add_alignment(algorithm, sample, out)
    out = _maybe_add_disambiguate(algorithm, sample, out)
    out = _maybe_add_counts(algorithm, sample, out)
    out = _maybe_add_cufflinks(algorithm, sample, out)
    out = _maybe_add_oncofuse(algorithm, sample, out)
    out = _maybe_add_rnaseq_variant_file(algorithm, sample, out)
    return _add_meta(out, sample)

def _get_files_chipseq(sample):
    out = []
    algorithm = sample["config"]["algorithm"]
    out = _maybe_add_summary(algorithm, sample, out)
    out = _maybe_add_alignment(algorithm, sample, out)
    return _add_meta(out, sample)

def _add_meta(xs, sample=None, config=None):
    out = []
    for x in xs:
        x["mtime"] = shared.get_file_timestamp(x["path"])
        if sample and "sample" not in x:
            if isinstance(sample["name"], (tuple, list)):
                name = sample["name"][-1]
            else:
                name = "%s-%s" % (sample["name"],
                                  run_info.clean_name(sample["description"]))
            x["sample"] = name
        if config:
            if "fc_name" in config and "fc_date" in config:
                x["run"] = "%s_%s" % (config["fc_date"], config["fc_name"])
            else:
                x["run"] = "project_%s" % datetime.datetime.now().strftime("%Y-%m-%d")
        out.append(x)
    return out

def _get_files_variantcall(sample):
    """Return output files for the variant calling pipeline.
    """
    out = []
    algorithm = sample["config"]["algorithm"]
    out = _maybe_add_summary(algorithm, sample, out)
    out = _maybe_add_alignment(algorithm, sample, out)
    out = _maybe_add_disambiguate(algorithm, sample, out)
    out = _maybe_add_variant_file(algorithm, sample, out)
    out = _maybe_add_sv(algorithm, sample, out)
    out = _maybe_add_validate(algorithm, sample, out)
    return _add_meta(out, sample)

def _maybe_add_validate(algorith, sample, out):
    for i, plot in enumerate(tz.get_in(("validate", "grading_plots"), sample, [])):
        ptype = os.path.splitext(plot)[-1].replace(".", "")
        out.append({"path": plot,
                    "type": ptype,
                    "ext": "validate%s" % ("" if i == 0 else "-%s" % (i + 1))})
    return out

def _maybe_add_rnaseq_variant_file(algorithm, sample, out):
    if sample.get("vrn_file"):
        out.append({"path": sample.get("vrn_file"),
                    "type": "vcf",
                    "ext": "vcf"})
    return out

def _maybe_add_variant_file(algorithm, sample, out):
    if sample.get("align_bam") is not None and sample.get("vrn_file"):
        for x in sample["variants"]:
            out.extend(_get_variant_file(x, ("vrn_file",)))
            if x.get("bed_file"):
                out.append({"path": x["bed_file"],
                            "type": "bed",
                            "ext": "%s-callregions" % x["variantcaller"],
                            "variantcaller": x["variantcaller"]})
            if x.get("vrn_stats"):
                for extra, fname in x["vrn_stats"].items():
                    ext = utils.splitext_plus(fname)[-1].replace(".", "")
                    out.append({"path": fname,
                                "type": ext,
                                "ext": "%s-%s" % (x["variantcaller"], extra),
                                "variantcaller": x["variantcaller"]})
    return out

def _maybe_add_sv(algorithm, sample, out):
    if sample.get("align_bam") is not None and sample.get("sv"):
        for svcall in sample["sv"]:
            out.extend(_get_variant_file(svcall, ("vrn_file",)))
            out.extend(_get_variant_file(svcall, ("bedpe_file",)))
            out.extend(_get_variant_file(svcall, ("cnr",)))
            out.extend(_get_variant_file(svcall, ("cns",)))
            out.extend(_get_variant_file(svcall, ("cnr_bed",)))
            out.extend(_get_variant_file(svcall, ("cnr_bedgraph",)))
            if "sample_bed" in svcall:
                out.append({"path": svcall["sample_bed"],
                            "type": "bed",
                            "ext": "%s-sample" % svcall["variantcaller"],
                            "variantcaller": svcall["variantcaller"]})
            if "validate" in svcall:
                for vkey in ["csv", "plot", "df"]:
                    vfile = tz.get_in(["validate", vkey], svcall)
                    if vfile:
                        to_u = []
                        if isinstance(vfile, dict):
                            for svtype, fname in vfile.items():
                                to_u.append((fname, "-%s" % svtype))
                        else:
                            to_u.append((vfile, "-%s" % vkey if vkey in ["df"] else ""))
                        for vfile, ext in to_u:
                            vext = os.path.splitext(vfile)[-1].replace(".", "")
                            out.append({"path": vfile,
                                        "type": vext,
                                        "ext": "%s-sv-validate%s" % (svcall["variantcaller"], ext),
                                        "variantcaller": svcall["variantcaller"]})
            if "plot" in svcall:
                for plot_name, fname in svcall["plot"].items():
                    ext = os.path.splitext(fname)[-1].replace(".", "")
                    out.append({"path": fname,
                                "type": ext,
                                "ext": "%s-%s" % (svcall["variantcaller"], plot_name),
                                "variantcaller": svcall["variantcaller"]})
    return out

def _get_variant_file(x, key):
    """Retrieve VCF file with the given key if it exists, handling bgzipped.
    """
    out = []
    fname = utils.get_in(x, key)
    upload_key = list(key)
    upload_key[-1] = "do_upload"
    do_upload = tz.get_in(tuple(upload_key), x, True)
    if fname and do_upload:
        if fname.endswith(".vcf.gz"):
            out.append({"path": fname,
                        "type": "vcf.gz",
                        "ext": x["variantcaller"],
                        "variantcaller": x["variantcaller"]})
            if utils.file_exists(fname + ".tbi"):
                out.append({"path": fname + ".tbi",
                            "type": "vcf.gz.tbi",
                            "index": True,
                            "ext": x["variantcaller"],
                            "variantcaller": x["variantcaller"]})
        elif fname.endswith((".vcf", ".bed", ".bedpe", ".bedgraph", ".cnr", ".cns", ".cnn")):
            ftype = utils.splitext_plus(fname)[-1][1:]
            out.append({"path": fname,
                        "type": ftype,
                        "ext": x["variantcaller"],
                        "variantcaller": x["variantcaller"]})
    return out

def _maybe_add_summary(algorithm, sample, out):
    out = []
    if "summary" in sample:
{'fraction_non_alphanumeric': 0.13818679751240942, 'fraction_numerical': 0.0015975352313573343, 'mean_word_length': 2.003427004797807, 'pattern_counts': {'":': 90, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3480727', 'n_tokens_mistral': 5256, 'n_tokens_neox': 5173, 'n_words': 1185}