<commit_before>"""
Common functions for analysis jobs.

"""
import asyncio
import json
import logging
import os
import pathlib
import shutil
from typing import Union

import pymongo.errors

import virtool.caches.db
import virtool.caches.utils
import virtool.history.db
import virtool.jobs.fastqc
import virtool.jobs.job
import virtool.jobs.utils
import virtool.samples.db
import virtool.samples.utils
import virtool.subtractions.utils
import virtool.utils
from virtool.jobs.caching import create_cache, set_cache_id

TRIMMING_PROGRAM = "skewer-0.2.2"

logger = logging.getLogger(__name__)


async def check_db(job: virtool.jobs.job.Job):
    """
    Get some initial information from the database that will be required during the course of the job. Store the data
    in `job.params`.

    Returned values are used by all analysis job types.

    TODO: move to virtool-workflow and make specific to analysis workflows

    :param job: the job object
    :return:
    """
    logger.debug("Retrieving job parameters")

    job.params = {
        # The document id for the sample being analyzed. and the analysis document the results will be committed to.
        "sample_id": job.task_args["sample_id"],

        # The document id for the reference to analyze against.
        "ref_id": job.task_args["ref_id"],

        # The document id for the analysis being run.
        "analysis_id": job.task_args["analysis_id"],
    }

    # The document for the sample being analyzed. Assigned after database connection is made.
    sample = await job.db.samples.find_one(job.params["sample_id"])

    # The parent folder for all data associated with the sample
    sample_path = os.path.join(job.settings["data_path"], "samples", job.params["sample_id"])

    analysis_path = os.path.join(sample_path, "analysis", job.params["analysis_id"])

    analysis = await job.db.analyses.find_one(job.params["analysis_id"], ["subtraction"])

    temp_analysis_path = os.path.join(job.temp_dir.name, job.task_args["analysis_id"])

    job.params.update({
        # The path to the directory where all analysis result files will be written.
        "analysis_path": analysis_path,
        "index_path": os.path.join(
            job.settings["data_path"],
            "references",
            job.params["ref_id"],
            job.task_args["index_id"],
            "reference"
        ),
        "sample_path": sample_path,
        "paired": sample["paired"],
        #: The number of reads in the sample library. Assigned after database connection is made.
        "read_count": int(sample["quality"]["count"]),
        "sample_read_length": int(sample["quality"]["length"][1]),
        "library_type": sample["library_type"],
        "reads_path": os.path.join(job.temp_dir.name, "reads"),
        "subtraction_path": virtool.subtractions.utils.join_subtraction_index_path(job.settings, analysis["subtraction"]["id"]),
        "raw_path": os.path.join(job.temp_dir.name, "raw"),
        "temp_cache_path": os.path.join(job.temp_dir.name, "cache"),
        "temp_analysis_path": temp_analysis_path
    })

    index_info = await get_index_info(
        job.db,
        job.settings,
        job.task_args["index_id"]
    )

    job.params.update(index_info)

    read_paths = [os.path.join(job.params["reads_path"], "reads_1.fq.gz")]

    if job.params["paired"]:
        read_paths.append(os.path.join(job.params["reads_path"], "reads_2.fq.gz"))

    job.params["read_paths"] = read_paths


async def make_analysis_dir(job):
    """
    Make a directory for the analysis in the sample/analysis directory.

    TODO: integrate into analysis working directory tree creation at start of all analysis workflows in virtool-workflow

    """
    await job.run_in_executor(os.mkdir, job.params["temp_analysis_path"])
    await job.run_in_executor(os.mkdir, job.params["raw_path"])
    await job.run_in_executor(os.mkdir, job.params["reads_path"])


async def prepare_reads(job: virtool.jobs.job.Job):
    """
    The first step of all analysis jobs.

    Prepares the sample read data by either:
    - fetching an existing trim cache if the parameters match
    - creating a new trim cache if one doesn't exist
    - fetching legacy sample data (to be deprecated)

    TODO: integrate into all analysis workflows in virtool-workflow.

    :param job: the job object

    """
    paired = job.params["paired"]

    parameters = get_trimming_parameters(
        paired,
        job.params["library_type"],
        job.params["sample_read_length"]
    )

    cache = await virtool.caches.db.find_and_wait(<commit_msg>Use asyncio.gather() when making analysis directory<commit_after>"""
Common functions for analysis jobs.

"""
import asyncio
import json
import logging
import os
import pathlib
import shutil
from asyncio import gather
from typing import Union

import pymongo.errors

import virtool.caches.db
import virtool.caches.utils
import virtool.history.db
import virtool.jobs.fastqc
import virtool.jobs.job
import virtool.jobs.utils
import virtool.samples.db
import virtool.samples.utils
import virtool.subtractions.utils
import virtool.utils
from virtool.jobs.caching import create_cache, set_cache_id

TRIMMING_PROGRAM = "skewer-0.2.2"

logger = logging.getLogger(__name__)


async def check_db(job: virtool.jobs.job.Job):
    """
    Get some initial information from the database that will be required during the course of the job. Store the data
    in `job.params`.

    Returned values are used by all analysis job types.

    TODO: move to virtool-workflow and make specific to analysis workflows

    :param job: the job object
    :return:
    """
    logger.debug("Retrieving job parameters")

    job.params = {
        # The document id for the sample being analyzed. and the analysis document the results will be committed to.
        "sample_id": job.task_args["sample_id"],

        # The document id for the reference to analyze against.
        "ref_id": job.task_args["ref_id"],

        # The document id for the analysis being run.
        "analysis_id": job.task_args["analysis_id"],
    }

    # The document for the sample being analyzed. Assigned after database connection is made.
    sample = await job.db.samples.find_one(job.params["sample_id"])

    # The parent folder for all data associated with the sample
    sample_path = os.path.join(job.settings["data_path"], "samples", job.params["sample_id"])

    analysis_path = os.path.join(sample_path, "analysis", job.params["analysis_id"])

    analysis = await job.db.analyses.find_one(job.params["analysis_id"], ["subtraction"])

    temp_analysis_path = os.path.join(job.temp_dir.name, job.task_args["analysis_id"])

    job.params.update({
        # The path to the directory where all analysis result files will be written.
        "analysis_path": analysis_path,
        "index_path": os.path.join(
            job.settings["data_path"],
            "references",
            job.params["ref_id"],
            job.task_args["index_id"],
            "reference"
        ),
        "sample_path": sample_path,
        "paired": sample["paired"],
        #: The number of reads in the sample library. Assigned after database connection is made.
        "read_count": int(sample["quality"]["count"]),
        "sample_read_length": int(sample["quality"]["length"][1]),
        "library_type": sample["library_type"],
        "reads_path": os.path.join(job.temp_dir.name, "reads"),
        "subtraction_path": virtool.subtractions.utils.join_subtraction_index_path(job.settings, analysis["subtraction"]["id"]),
        "raw_path": os.path.join(job.temp_dir.name, "raw"),
        "temp_cache_path": os.path.join(job.temp_dir.name, "cache"),
        "temp_analysis_path": temp_analysis_path
    })

    index_info = await get_index_info(
        job.db,
        job.settings,
        job.task_args["index_id"]
    )

    job.params.update(index_info)

    read_paths = [os.path.join(job.params["reads_path"], "reads_1.fq.gz")]

    if job.params["paired"]:
        read_paths.append(os.path.join(job.params["reads_path"], "reads_2.fq.gz"))

    job.params["read_paths"] = read_paths


async def make_analysis_dir(job):
    """
    Make a directory for the analysis in the sample/analysis directory.

    TODO: integrate into analysis working directory tree creation at start of all analysis workflows in virtool-workflow

    """
    await gather(
        job.run_in_executor(os.mkdir, job.params["temp_analysis_path"]),
        job.run_in_executor(os.mkdir, job.params["raw_path"]),
        job.run_in_executor(os.mkdir, job.params["reads_path"])
    )


async def prepare_reads(job: virtool.jobs.job.Job):
    """
    The first step of all analysis jobs.

    Prepares the sample read data by either:
    - fetching an existing trim cache if the parameters match
    - creating a new trim cache if one doesn't exist
    - fetching legacy sample data (to be deprecated)

    TODO: integrate into all analysis workflows in virtool-workflow.

    :param job: the job object

    """
    paired = job.params["paired"]

    parameters = get_trimming_parameters(
        paired,
        job.params["library_type"],
        job.params["sample_read_length"]
    )

    cache = await virtool.caches.db.find_and_wait(
{'fraction_non_alphanumeric': 0.10928902794260438, 'fraction_numerical': 0.0012946380407810982, 'mean_word_length': 3.3582510578279265, 'pattern_counts': {'":': 30, '<': 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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2726818', 'n_tokens_mistral': 2843, 'n_tokens_neox': 2740, 'n_words': 749}