#!/usr/bin/env python3

"""
Purpose
-------

This module is intended execute megahit on paired-end FastQ files.

Expected input
--------------

The following variables are expected whether using NextFlow or the
:py:func:`main` executor.

- ``sample_id`` : Sample Identification string.
    - e.g.: ``'SampleA'``
- ``fastq_pair`` : Pair of FastQ file paths.
    - e.g.: ``'SampleA_1.fastq.gz SampleA_2.fastq.gz'``
- ``kmers`` : Setting for megahit kmers. Can be either ``'auto'``, \
    ``'default'`` or a user provided list. All must be odd, in the range 15-255, increment <= 28
    - e.g.: ``'auto'`` or ``'default'`` or ``'55 77 99 113 127'``
- ``clear`` : If 'true', remove the input fastq files at the end of the
    component run, IF THE FILES ARE IN THE WORK DIRECTORY

Generated output
----------------

- ``contigs.fa`` : Main output of megahit with the assembly
    - e.g.: ``contigs.fa``
- ``megahit_status`` :  Stores the status of the megahit run. If it was \
    successfully executed, it stores ``'pass'``. Otherwise, it stores the\
    ``STDERR`` message.
    - e.g.: ``'pass'``

Code documentation
------------------

"""

__version__ = "1.0.1"
__build__ = "26042018"
__template__ = "megahit-nf"

import os
import re
import subprocess

from subprocess import PIPE

from flowcraft_utils.flowcraft_base import get_logger, MainWrapper

logger = get_logger(__file__)

def is_odd(k_mer):
    for i in k_mer:
        if i % 2 != 0:
            return True
    return False

def __get_version_megahit():

    try:

        cli = ["megahit", "--version"]
        p = subprocess.Popen(cli, stdout=PIPE, stderr=PIPE)
        stdout, _ = p.communicate()

        version = stdout.strip().split()[-1][1:].decode("utf8")

    except Exception as e:
        logger.debug(e)
        version = "undefined"

    return {
        "program": "megahit",
        "version": version,
    }

if __file__.endswith(".command.sh"):
    SAMPLE_ID = '$sample_id'
    FASTQ_PAIR = '$fastq_pair'.split()
    MAX_LEN = int('$max_len'.strip())
    KMERS = '$kmers'.strip()
    MEM = '$task.memory'
    CLEAR = '$clear'
    logger.debug("Running {} with parameters:".format(
        os.path.basename(__file__)))
    logger.debug("SAMPLE_ID: {}".format(SAMPLE_ID))
    logger.debug("FASTQ_PAIR: {}".format(FASTQ_PAIR))
    logger.debug("MAX_LEN: {}".format(MAX_LEN))
    logger.debug("KMERS: {}".format(KMERS))
    logger.debug("CLEAR: {}".format(CLEAR))

def set_kmers(kmer_opt, max_read_len):
    """Returns a kmer list based on the provided kmer option and max read len.

    Parameters
    ----------
    kmer_opt : str
        The k-mer option. Can be either ``'auto'``, ``'default'`` or a
        sequence of space separated integers, ``'23, 45, 67'``.
    max_read_len : int
        The maximum read length of the current sample.

    Returns
    -------
    kmers : list
        List of k-mer values that will be provided to megahit.

    """

    logger.debug("Kmer option set to: {}".format(kmer_opt))

    # Check if kmer option is set to auto
    if kmer_opt == "auto":

        if max_read_len >= 175:
            kmers = [55, 77, 99, 113, 127]
        else:
            kmers = [21, 33, 55, 67, 77]

        logger.debug("Kmer range automatically selected based on max read"
                     "length of {}: {}".format(max_read_len, kmers))

    # Check if manual k-mers were specified
    elif len(kmer_opt.split()) > 1:

        kmers = kmer_opt.split()
        if kmers[0]<15 or kmers[-1]>255 or is_odd(kmers):
            kmers = []
            logger.debug("Kmer out of range or with even numbers"
                         "(will be automatically determined by megahit")
        else:
            logger.debug("Kmer range manually set to: {}".format(kmers))

    else:

        kmers = []
        logger.debug("Kmer range set to empty (will be automatically "
                     "determined by megahit")

    return kmers

def fix_contig_names(asseembly_path):
    """Removes whitespace from the assembly contig names

    Parameters
    ----------
    asseembly_path : path to assembly file

    Returns
    -------
    str:
        Path to new assembly file with fixed contig names
    """

    fixed_assembly = "fixed_assembly.fa"

    with open(asseembly_path) as in_hf, open(fixed_assembly, "w") as ou_fh:

        for line in in_hf:

            if line.startswith(">"):
                fixed_line = line.replace(" ", "_")
                ou_fh.write(fixed_line)
            else:
                ou_fh.write(line)

    return fixed_assembly

def clean_up(fastq):
    """
    Cleans the temporary fastq files. If they are symlinks, the link
    source is removed

    Parameters
    ----------
    fastq : list
        List of fastq files.
    """

    for fq in fastq:
        # Get real path of fastq files, following symlinks
        rp = os.path.realpath(fq)
        logger.debug("Removing temporary fastq file path: {}".format(rp))
        if re.match(".*/work/.{2}/.{30}/.*", rp):
            os.remove(rp)

@MainWrapper
def main(sample_id, fastq_pair, max_len, kmer, mem, clear):
    """Main executor of the megahit template.

    Parameters
    ----------
    sample_id : str
        Sample Identification string.
    fastq_pair : list
        Two element list containing the paired FastQ files.
    max_len : int
        Maximum read length. This value is determined in
        :py:class:`templates.integrity_coverage`
    kmer : str
        Can be either ``'auto'``, ``'default'`` or a
        sequence of space separated integers, ``'23, 45, 67'``.

    """

    logger.info("Starting megahit")

    logger.info("Setting megahit kmers")
    kmers = set_kmers(kmer, max_len)
    logger.info("megahit kmers set to: {}".format(kmers))

    mem_bytes = int(mem.replace(" GB", "")) * 1073741824

    cli = [
        "megahit",
        "--num-cpu-threads",
        "$task.cpus",
        "--memory",
        str(mem_bytes),
        "-o",
        "megahit"
    ]

    # Add kmers, if any were specified
    if kmers:
        cli += [
            "--k-list",
            ",".join([str(x) for x in kmers])
        ]

    # Add FastQ files
    cli += [
        "-1",
        fastq_pair[0],
        "-2",
        fastq_pair[1]
    ]

    logger.debug("Running megahit subprocess with command: {}".format(cli))

    p = subprocess.Popen(cli, stdout=PIPE, stderr=PIPE)
    stdout, stderr = p.communicate()

    # Attempt to decode STDERR output from bytes. If unsuccessful, coerce to
    # string
    try:
        stderr = stderr.decode("utf8")
        stdout = stdout.decode("utf8")
    except (UnicodeDecodeError, AttributeError):
        stderr = str(stderr)
        stdout = str(stdout)

    logger.info("Finished megahit subprocess with STDOUT:\\n"
                "======================================\\n{}".format(stdout))
    logger.info("Fished megahit subprocesswith STDERR:\\n"
                "======================================\\n{}".format(stderr))
    logger.info("Finished megahit with return code: {}".format(
        p.returncode))

    with open(".status", "w") as fh:
        if p.returncode != 0:
            fh.write("error")
            return
        else:
            fh.write("pass")

    assembly_path = "megahit/final.contigs.fa"
    fixed_assembly = fix_contig_names(assembly_path)

    # Change the default final.contigs.fa assembly name to a more informative
    #  one
    if "_trim." in fastq_pair[0]:
        sample_id += "_trim"
    # Get megahit version for output name
    info = __get_version_megahit()

    assembly_file = "{}_megahit{}.fasta".format(
        sample_id, info["version"].replace(".", ""))
    os.rename(fixed_assembly, assembly_file)
    logger.info("Setting main assembly file to: {}".format(assembly_file))

    # Remove input fastq files when clear option is specified.
    # Only remove temporary input when the expected output exists.
    if clear == "true" and os.path.exists(assembly_file):
        clean_up(fastq_pair)

if __name__ == '__main__':

    main(SAMPLE_ID, FASTQ_PAIR, MAX_LEN, KMERS, MEM, CLEAR)


#!/usr/bin/env python3

"""
Purpose
-------

This module is intended parse the results of the Abricate for one or more
samples.

Expected input
--------------

The following variables are expected whether using NextFlow or the
:py:func:`main` executor.

- ``abricate_files`` : Path to abricate output file.
    - e.g.: ``'abr_resfinder.tsv'``

Generated output
----------------

None

Code documentation
------------------

"""

__version__ = "1.0.1"
__build__ = "26032018"
__template__ = "process_abricate-nf"

import re
import os
import json
import operator
import subprocess

from subprocess import PIPE

from flowcraft_utils.flowcraft_base import get_logger, MainWrapper

logger = get_logger(__file__)

def __get_version_abricate():

    try:

        # Get abricate version
        cli = ["abricate", "--version"]
        p = subprocess.Popen(cli, stdout=PIPE, stderr=PIPE)
        stdout, _ = p.communicate()

        version = stdout.strip().split()[-1].decode("utf8")

    except Exception as e:
        logger.debug(e)
        version = "undefined"

    try:

        # Get abricate database versions
        cli = ["abricate", "--list"]
        p = subprocess.Popen(cli, stdout=PIPE, stderr=PIPE)
        dbout, _ = p.communicate()

        databases = [[u.decode("utf8") for u in i.strip().split()]
                     for i in dbout.splitlines()][1:]

    except Exception as e:
        logger.debug(e)
        databases = "undefined"

    return {
        "program": "abricate",
        "version": version,
        "databases": databases
    }

if __file__.endswith(".command.sh"):
    ABRICATE_FILES = '$abricate_file'.split()
    logger.debug("Running {} with parameters:".format(
        os.path.basename(__file__)))
    logger.debug("ABRICATE_FILE: {}".format(ABRICATE_FILES))

class Abricate:
    """Main parser for Abricate output files.

    This class parses one or more output files from Abricate, usually from
    different databases. In addition to the parsing methods, it also provides
    a flexible method to filter and re-format the content of the abricate
    files.

    Parameters
    ----------
    fls : list
       List of paths to Abricate output files.
    """

    def __init__(self, fls):

        self.storage = {}
        """
        dic: Main storage of Abricate's file content. Each entry corresponds
        to a single line and contains the keys::

            - ``log_file``: Name of the summary log file containing abricate
              results
            - ``infile``: Input file of Abricate.
            - ``reference``: Reference of the query sequence.
            - ``seq_range``: Range of the query sequence in the database
             sequence.
            - ``gene``: AMR gene name.
            - ``accession``: The genomic source of the sequence.
            - ``database``: The database the sequence came from.
            - ``coverage``: Proportion of gene covered.
            - ``identity``: Proportion of exact nucleotide matches.
        """

        self._key = 0
        """
        int: Arbitrary key for unique entries in the storage attribute
        """

        self.parse_files(fls)

    def parse_files(self, fls):
        """Public method for parsing abricate output files.

        This method is called at at class instantiation for the provided
        output files. Additional abricate output files can be added using
        this method after the class instantiation.

        Parameters
        ----------
        fls : list
            List of paths to Abricate files

        """

        for f in fls:
            # Make sure paths exists
            if os.path.exists(f):
                self._parser(f)
            else:
                logger.warning("File {} does not exist".format(f))

    def _parser(self, fl):
        """Parser for a single abricate output file.

        This parser will scan a single Abricate output file and populate
        the :py:attr:`Abricate.storage` attribute.

        Parameters
        ----------
        fl : str
            Path to abricate output file

        Notes
        -----
        This method will populate the :py:attr:`Abricate.storage` attribute
        with all compliant lines in the abricate output file. Entries are
        inserted using an arbitrary key that is set by the
        :py:attr:`Abricate._key` attribute.

        """

        with open(fl) as fh:

            for line in fh:
                # Skip header and comment lines
                if line.startswith("#") or line.strip() == "":
                    continue

                fields = line.strip().split("\t")

                try:
                    coverage = float(fields[8])
                except ValueError:
                    coverage = None
                try:
                    identity = float(fields[9])
                except ValueError:
                    identity = None

                try:
                    accession = fields[11]
                except IndexError:
                    accession = None

                self.storage[self._key] = {
                    "log_file": os.path.basename(fl),
                    "infile": fields[0],
                    "reference": fields[1],
                    "seq_range": (int(fields[2]), int(fields[3])),
                    "gene": fields[4],
                    "accession": accession,
                    "database": fields[10],
                    "coverage": coverage,
                    "identity": identity
                }

                self._key += 1

    @staticmethod
    def _test_truth(x, op, y):
        """ Test the truth of a comparison between x and y using an operator.

        If you want to compare '100 > 200', this method can be called as
        self._test_truth(100, ">", 200).

        Parameters
        ----------
        x : int
            Arbitrary value to compare in the left.
        op : str
            Comparison operator.
        y : int
            Arbitrary value to compare in the right.

        Returns
        -------
        x : bool
            The 'truthness' of the test.
        """

        ops = {
            ">": operator.gt,
            "<": operator.lt,
            ">=": operator.ge,
            "<=": operator.le,
            "==": operator.eq,
            "!=": operator.ne
        }

        return ops[op](x, y)

    def iter_filter(self, filters, databases=None, fields=None,
                    filter_behavior="and"):
        """General purpose filter iterator.

        This general filter iterator allows the filtering of entries based
        on one or more custom filters. These filters must contain
        an entry of the `storage` attribute, a comparison operator, and the
        test value. For example, to filter out entries with coverage below 80::

            my_filter = ["coverage", ">=", 80]

        Filters should always be provide as a list of lists::

            iter_filter([["coverage", ">=", 80]])
            # or
            my_filters = [["coverage", ">=", 80],
                          ["identity", ">=", 50]]

            iter_filter(my_filters)

        As a convenience, a list of the desired databases can be directly
        specified using the `database` argument, which will only report
        entries for the specified databases::

            iter_filter(my_filters, databases=["plasmidfinder"])

        By default, this method will yield the complete entry record. However,
        the returned filters can be specified using the `fields` option::

            iter_filter(my_filters, fields=["reference", "coverage"])

        Parameters
        ----------
        filters : list
            List of lists with the custom filter. Each list should have three
            elements. (1) the key from the entry to be compared; (2) the
            comparison operator; (3) the test value. Example:
                ``[["identity", ">", 80]]``.
        databases : list
            List of databases that should be reported.
        fields : list
            List of fields from each individual entry that are yielded.
        filter_behavior : str
            options: ``'and'`` ``'or'``
            Sets the behaviour of the filters, if multiple filters have been
            provided. By default it is set to ``'and'``, which means that an
            entry has to pass all filters. It can be set to ``'or'``, in which
            case one one of the filters has to pass.

        yields
        ------
        dic : dict
            Dictionary object containing a :py:attr:`Abricate.storage` entry
            that passed the filters.

        """

        if filter_behavior not in ["and", "or"]:
            raise ValueError("Filter behavior must be either 'and' or 'or'")

        for dic in self.storage.values():

            # This attribute will determine whether an entry will be yielded
            # or not
            _pass = False

            # Stores the flags with the test results for each filter
            # The results will be either True or False
            flag = []

            # Filter for databases
            if databases:
                # Skip entry if not in specified database
                if dic["database"] not in databases:
                    continue

            # Apply filters
            for f in filters:
                # Get value of current filter
                val = dic[f[0]]
                if not self._test_truth(val, f[1], f[2]):
                    flag.append(False)
                else:
                    flag.append(True)

            # Test whether the entry will pass based on the test results
            # and the filter behaviour
            if filter_behavior == "and":
                if all(flag):
                    _pass = True
            elif filter_behavior == "or":
                if any(flag):
                    _pass = True

            if _pass:
                if fields:
                    yield dict((x, y) for x, y in dic.items() if x in fields)
                else:
                    yield dic

    def get_filter(self, *args, **kwargs):
        """ Wrapper of the iter_filter method that returns a list with results

        It should be called exactly as in the `iter_filter`

        Returns
        -------
        _ : list
            List of dictionary entries that passed the filters in the
            `iter_filter` method.

        See Also
        --------
        iter_filter
        """

        return list(self.iter_filter(*args, **kwargs))

class AbricateReport(Abricate):
    """Report generator for single Abricate output files

    This class is intended to parse an Abricate output file from a single
    sample and database and generates a JSON report for the report webpage.

    Parameters
    ----------
    fls : list
       List of paths to Abricate output files.
    database : (optional) str
        Name of the database for the current report. If not provided, it will
        be inferred based on the first entry of the Abricate file.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    @staticmethod
    def _get_contig_id(contig_str):
        """Tries to retrieve contig id. Returns the original string if it
        is unable to retrieve the id.

        Parameters
        ----------
        contig_str : str
            Full contig string (fasta header)

        Returns
        -------
        str
            Contig id
        """

        contig_id = contig_str

        try:
            contig_id = re.search(".*NODE_([0-9]*)_.*", contig_str).group(1)
        except AttributeError:
            pass

        try:
            contig_id = re.search(".*Contig_([0-9]*)_.*", contig_str).group(1)
        except AttributeError:
            pass

        return contig_id

    def get_plot_data(self):
        """ Generates the JSON report to plot the gene boxes

        Following the convention of the reports platform, this method returns
        a list of JSON/dict objects with the information about each entry in
        the abricate file. The information contained in this JSON is::

            {contig_id: <str>,
             seqRange: [<int>, <int>],
             gene: <str>,
             accession: <str>,
             coverage: <float>,
             identity: <float>
             }

        Note that the `seqRange` entry contains the position in the
        corresponding contig, not the absolute position in the whole assembly.

        Returns
        -------
        json_dic : list
            List of JSON/dict objects with the report data.
        """

        json_dic = {"plotData": []}
        sample_dic = {}
        sample_assembly_map = {}

        for entry in self.storage.values():

            sample_id = re.match("(.*)_abr", entry["log_file"]).groups()[0]
            if sample_id not in sample_dic:
                sample_dic[sample_id] = {}

            # Get contig ID using the same regex as in `assembly_report.py`
            # template
            contig_id = self._get_contig_id(entry["reference"])
            # Get database
            database = entry["database"]
            if database not in sample_dic[sample_id]:
                sample_dic[sample_id][database] = []

            # Update the sample-assembly correspondence dict
            if sample_id not in sample_assembly_map:
                sample_assembly_map[sample_id] = entry["infile"]

            sample_dic[sample_id][database].append(
                {"contig": contig_id,
                 "seqRange": entry["seq_range"],
                 "gene": entry["gene"].replace("'", ""),
                 "accession": entry["accession"],
                 "coverage": entry["coverage"],
                 "identity": entry["identity"],
                 },
            )

        for sample, data in sample_dic.items():
            json_dic["plotData"].append(
                {
                    "sample": sample,
                    "data": {"abricateXrange": data},
                    "assemblyFile": sample_assembly_map[sample]
                }
            )

        return json_dic

    def get_table_data(self):
        """

        Returns
        -------

        """

        gene_storage = {}
        json_dic = {"tableRow": []}
        logger.info("Generating JSON table data")

        # Collect the gene lists for each database
        for key, entry in self.storage.items():

            # Retrieve and initiate new sample entry, if not present already
            logger.debug("Retrieving sample if from: {}".format(
                entry["infile"]))
            sample_id = re.match("(.*)_abr", entry["log_file"]).groups()[0]
            database = entry["database"]

            if sample_id not in gene_storage:
                gene_storage[sample_id] = {}

            if database not in gene_storage[sample_id]:
                gene_storage[sample_id][database] = []

            gene_storage[sample_id][database].append(
                entry["gene"].replace("'", "").replace('"', '')
            )

        # For each database, create the JSON report
        for sample, table_data in gene_storage.items():

            json_dic["tableRow"].append({
                "sample": sample,
                "data": []
            })

            for db, gene_list in table_data.items():

                ind_json = {
                    "table": "abricate",
                    "header": db,
                    "value": len(gene_list),
                    "geneList": gene_list
                }
                json_dic["tableRow"][-1]["data"].append(ind_json)

        return json_dic

    def write_report_data(self):
        """Writes the JSON report to a json file
        """

        json_plot = self.get_plot_data()
        json_table = self.get_table_data()

        json_dic = {**json_plot, **json_table}

        with open(".report.json", "w") as json_report:
            json_report.write(json.dumps(json_dic, separators=(",", ":")))

if __name__ == '__main__':

    @MainWrapper
    def main(abr_file):

        abr = AbricateReport(fls=abr_file)
        abr.write_report_data()

    main(ABRICATE_FILES)


#!/usr/bin/env python3

"""
Purpose
-------

This module is intended to provide a summary report for a given assembly
in Fasta format.

Expected input
--------------

The following variables are expected whether using NextFlow or the
:py:func:`main` executor.

- ``sample_id`` : Sample Identification string.
    - e.g.: ``'SampleA'``
- ``assembly`` : Path to assembly file in Fasta format.
    - e.g.: ``'assembly.fasta'``

Generated output
----------------

- ``${sample_id}_assembly_report.csv`` : CSV with summary information of the \
    assembly.
    - e.g.: ``'SampleA_assembly_report.csv'``

Code documentation
------------------

"""

__version__ = "1.0.1"
__build__ = "16012018"
__template__ = "assembly_report-nf"

import os
import re
import json
import traceback
import subprocess

from collections import OrderedDict
from subprocess import PIPE

from flowcraft_utils.flowcraft_base import get_logger, MainWrapper

logger = get_logger(__file__)

def __get_version_pilon():

    pilon_path = "/NGStools/pilon-1.22.jar"

    try:

        cli = ["java", "-jar", pilon_path, "--version"]
        p = subprocess.Popen(cli, stdout=PIPE, stderr=PIPE)
        stdout, _ = p.communicate()

        version = stdout.split()[2].decode("utf8")

    except Exception as e:
        logger.debug(e)
        version = "undefined"

    return {
        "program": "Pilon",
        "version": version,
    }

if __file__.endswith(".command.sh"):
    SAMPLE_ID = '$sample_id'
    ASSEMBLY_FILE = '$assembly'
    COVERAGE_BP_FILE = '$coverage_bp'
    logger.debug("Running {} with parameters:".format(
        os.path.basename(__file__)))
    logger.debug("SAMPLE_ID: {}".format(SAMPLE_ID))
    logger.debug("ASSEMBLY_FILE: {}".format(ASSEMBLY_FILE))
    logger.debug("COVERAGE_BP_FILE: {}".format(COVERAGE_BP_FILE))

class Assembly:
    """Class that parses and filters an assembly file in Fasta format.

    This class parses an assembly file, collects a number
    of summary statistics and metadata from the contigs and reports.

    Parameters
    ----------
    assembly_file : str
        Path to assembly file.
    sample_id : str
        Name of the sample for the current assembly.
    """

    def __init__(self, assembly_file, sample_id):

        self.summary_info = OrderedDict([
            ("ncontigs", 0),
            ("avg_contig_size", []),
            ("n50", 0),
            ("total_len", 0),
            ("avg_gc", []),
            ("missing_data", 0)
        ])
        """
        OrderedDict: Initialize summary information dictionary. Contains keys:

            - ``ncontigs``: Number of contigs
            - ``avg_contig_size``: Average size of contigs
            - ``n50``: N50 metric
            - ``total_len``: Total assembly length
            - ``avg_gc``: Average GC proportion
            - ``missing_data``: Count of missing data characters
        """

        self.contigs = OrderedDict()
        """
        OrderedDict: Object that maps the contig headers to the corresponding
        sequence
        """

        self.contig_coverage = OrderedDict()
        """
        OrderedDict: Object that maps the contig headers to the corresponding
        list of per-base coverage
        """

        self.sample = sample_id
        """
        str: Sample id
        """

        self.contig_boundaries = {}
        """
        dict: Maps the boundaries of each contig in the genome
        """

        self._parse_assembly(assembly_file)

    def _parse_assembly(self, assembly_file):
        """Parse an assembly file in fasta format.

        This is a Fasta parsing method that populates the
        :py:attr:`Assembly.contigs` attribute with data for each contig in the
         assembly.

        Parameters
        ----------
        assembly_file : str
            Path to the assembly fasta file.

        """

        with open(assembly_file) as fh:

            header = None
            logger.debug("Starting iteration of assembly file: {}".format(
                assembly_file))

            for line in fh:

                # Skip empty lines
                if not line.strip():
                    continue

                if line.startswith(">"):
                    # Add contig header to contig dictionary
                    header = line[1:].strip()
                    self.contigs[header] = []

                else:
                    # Add sequence string for the current contig
                    self.contigs[header].append(line.strip())

            # After populating the contigs dictionary, convert the values
            # list into a string sequence
            self.contigs = OrderedDict(
                (header, "".join(seq)) for header, seq in self.contigs.items())

    @staticmethod
    def _get_contig_id(contig_str):
        """Tries to retrieve contig id. Returns the original string if it
        is unable to retrieve the id.

        Parameters
        ----------
        contig_str : str
            Full contig string (fasta header)

        Returns
        -------
        str
            Contig id
        """

        contig_id = contig_str

        try:
            contig_id = re.search(".*NODE_([0-9]*)_.*", contig_str).group(1)
        except AttributeError:
            pass

        try:
            contig_id = re.search(".*Contig_([0-9]*)_.*", contig_str).group(1)
        except AttributeError:
            pass

        return contig_id

    def get_summary_stats(self, output_csv=None):
        """Generates a CSV report with summary statistics about the assembly

        The calculated statistics are:

            - Number of contigs
            - Average contig size
            - N50
            - Total assembly length
            - Average GC content
            - Amount of missing data

        Parameters
        ----------
        output_csv: str
            Name of the output CSV file.
        """

        contig_size_list = []

        self.summary_info["ncontigs"] = len(self.contigs)

        for contig_id, sequence in self.contigs.items():

            logger.debug("Processing contig: {}".format(contig_id))

            # Get contig sequence size
            contig_len = len(sequence)

            # Add size for average contig size
            contig_size_list.append(contig_len)

            # Add to total assembly length
            self.summary_info["total_len"] += contig_len

            # Add to average gc
            self.summary_info["avg_gc"].append(
                sum(map(sequence.count, ["G", "C"])) / contig_len
            )

            # Add to missing data
            self.summary_info["missing_data"] += sequence.count("N")

        # Get average contig size
        logger.debug("Getting average contig size")
        self.summary_info["avg_contig_size"] = \
            sum(contig_size_list) / len(contig_size_list)

        # Get average gc content
        logger.debug("Getting average GC content")
        self.summary_info["avg_gc"] = \
            sum(self.summary_info["avg_gc"]) / len(self.summary_info["avg_gc"])

        # Get N50
        logger.debug("Getting N50")
        cum_size = 0
        for l in sorted(contig_size_list, reverse=True):
            cum_size += l
            if cum_size >= self.summary_info["total_len"] / 2:
                self.summary_info["n50"] = l
                break

        if output_csv:
            logger.debug("Writing report to csv")
            # Write summary info to CSV
            with open(output_csv, "w") as fh:
                summary_line = "{}, {}\\n".format(
                    self.sample, ",".join(
                        [str(x) for x in self.summary_info.values()]))
                fh.write(summary_line)

    def _get_window_labels(self, window):
        """Returns the mapping between sliding window points and their contigs,
        and the x-axis position of contig

        Parameters
        ----------
        window : int
            Size of the window.

        Returns
        -------
        xbars : list
            The x-axis position of the ending for each contig.
        labels : list
            The x-axis labels for each data point in the sliding window

        """

        # Get summary stats, if they have not yet been triggered
        if not self.summary_info:
            self.get_summary_stats()

        # Get contig boundary positon
        c = 0
        xbars = []
        for contig, seq in self.contigs.items():
            contig_id = self._get_contig_id(contig)
            self.contig_boundaries[contig_id] = [c, c + len(seq)]
            c += len(seq)
            xbars.append((contig_id, c, contig))

        return xbars

    @staticmethod
    def _gc_prop(s, length):
        """Get proportion of GC from a string

        Parameters
        ----------
        s : str
            Arbitrary string

        Returns
        -------
        x : float
            GC proportion.
        """

        gc = sum(map(s.count, ["c", "g"]))

        return gc / length

    def get_gc_sliding(self, window=2000):
        """Calculates a sliding window of the GC content for the assembly

        Returns
        -------
        gc_res : list
            List of GC proportion floats for each data point in the sliding
            window
        """

        gc_res = []

        # Get complete sequence to calculate sliding window values
        complete_seq = "".join(self.contigs.values()).lower()

        for i in range(0, len(complete_seq), window):

            seq_window = complete_seq[i:i + window]

            # Get GC proportion
            gc_res.append(round(self._gc_prop(seq_window, len(seq_window)), 2))

        return gc_res

    def _get_coverage_from_file(self, coverage_file):
        """

        Parameters
        ----------
        coverage_file

        Returns
        -------

        """

        with open(coverage_file) as fh:

            for line in fh:

                fields = line.strip().split()

                # Get header
                header = fields[0]
                coverage = int(fields[2])

                if header not in self.contig_coverage:
                    self.contig_coverage[header] = [coverage]
                else:
                    self.contig_coverage[header].append(coverage)

    def get_coverage_sliding(self, coverage_file, window=2000):
        """

        Parameters
        ----------
        coverage_file : str
            Path to file containing the coverage info at the per-base level
            (as generated by samtools depth)
        window : int
            Size of sliding window

        Returns
        -------

        """

        if not self.contig_coverage:
            self._get_coverage_from_file(coverage_file)

        # Stores the coverage results
        cov_res = []

        # Make flat list of coverage values across genome
        complete_cov = [x for y in self.contig_coverage.values() for x in y]

        for i in range(0, len(complete_cov), window):
            # Get coverage values for current window
            cov_window = complete_cov[i:i + window]
            # Get mean coverage
            cov_res.append(int(sum(cov_window) / len(cov_window)))

        return cov_res

@MainWrapper
def main(sample_id, assembly_file, coverage_bp_file=None):
    """Main executor of the assembly_report template.

    Parameters
    ----------
    sample_id : str
        Sample Identification string.
    assembly_file : str
        Path to assembly file in Fasta format.

    """

    logger.info("Starting assembly report")
    assembly_obj = Assembly(assembly_file, sample_id)

    logger.info("Retrieving summary statistics for assembly")
    assembly_obj.get_summary_stats("{}_assembly_report.csv".format(sample_id))

    size_dist = [len(x) for x in assembly_obj.contigs.values()]
    json_dic = {
        "tableRow": [{
            "sample": sample_id,
            "data": [
                {"header": "Contigs",
                 "value": assembly_obj.summary_info["ncontigs"],
                 "table": "assembly",
                 "columnBar": True},
                {"header": "Assembled BP",
                 "value": assembly_obj.summary_info["total_len"],
                 "table": "assembly",
                 "columnBar": True},
            ]
        }],
        "plotData": [{
            "sample": sample_id,
            "data": {
                "size_dist": size_dist
            }
        }]
    }

    if coverage_bp_file:
        try:
            window = 2000
            gc_sliding_data = assembly_obj.get_gc_sliding(window=window)
            cov_sliding_data = \
                assembly_obj.get_coverage_sliding(coverage_bp_file,
                                                  window=window)

            # Get total basepairs based on the individual coverage of each
            # contig bpx
            total_bp = sum(
                [sum(x) for x in assembly_obj.contig_coverage.values()]
            )

            # Add data to json report
            json_dic["plotData"][0]["data"]["genomeSliding"] = {
                "gcData": gc_sliding_data,
                "covData": cov_sliding_data,
                "window": window,
                "xbars": assembly_obj._get_window_labels(window),
                "assemblyFile": os.path.basename(assembly_file)
            }
            json_dic["plotData"][0]["data"]["sparkline"] = total_bp

        except:
            logger.error("Unexpected error creating sliding window data:\\n"
                         "{}".format(traceback.format_exc()))

    # Write json report
    with open(".report.json", "w") as json_report:

        json_report.write(json.dumps(json_dic, separators=(",", ":")))

    with open(".status", "w") as status_fh:
        status_fh.write("pass")

if __name__ == '__main__':

    main(SAMPLE_ID, ASSEMBLY_FILE, COVERAGE_BP_FILE)


#!/usr/bin/env python3

"""
Purpose
-------

This module is intended execute Skesa on paired-end FastQ files.

Expected input
--------------

The following variables are expected whether using NextFlow or the
:py:func:`main` executor.

- ``sample_id`` : Sample Identification string.
    - e.g.: ``'SampleA'``
- ``fastq_pair`` : Pair of FastQ file paths.
    - e.g.: ``'SampleA_1.fastq.gz SampleA_2.fastq.gz'``
- ``clear`` : If 'true', remove the input fastq files at the end of the
    component run, IF THE FILES ARE IN THE WORK DIRECTORY

Generated output
----------------

- ``${sample_id}_*.assembly.fasta`` : Main output of skesawith the assembly
    - e.g.: ``sample_1_skesa.fasta``
- ``clear`` : If 'true', remove the input fastq files at the end of the
    component run, IF THE FILES ARE IN THE WORK DIRECTORY

Code documentation
------------------

"""

__version__ = "1.0.2"
__build__ = "29062018"
__template__ = "skesa-nf"

import os
import re
import subprocess

from subprocess import PIPE

from flowcraft_utils.flowcraft_base import get_logger, MainWrapper

logger = get_logger(__file__)

def __get_version_skesa():

    try:

        cli = ["skesa", "--version"]
        p = subprocess.Popen(cli, stdout=PIPE, stderr=PIPE)
        _, err = p.communicate()

        try:
            version = re.search("v((\\..*))-", err.decode("utf8")).group(1)
        except AttributeError:
            version = "undefined"

    except Exception as e:
        logger.debug(e)
        version = "undefined"

    return {
        "program": "skesa",
        "version": version,
    }

if __file__.endswith(".command.sh"):
    SAMPLE_ID = '$sample_id'
    FASTQ_PAIR = '$fastq_pair'.split()
    CLEAR = '$clear'
    logger.debug("Running {} with parameters:".format(
        os.path.basename(__file__)))
    logger.debug("SAMPLE_ID: {}".format(SAMPLE_ID))
    logger.debug("FASTQ_PAIR: {}".format(FASTQ_PAIR))
    logger.debug("CLEAR: {}".format(CLEAR))

def clean_up(fastq):
    """
    Cleans the temporary fastq files. If they are symlinks, the link
    source is removed

    Parameters
    ----------
    fastq : list
        List of fastq files.
    """

    for fq in fastq:
        # Get real path of fastq files, following symlinks
        rp = os.path.realpath(fq)
        logger.debug("Removing temporary fastq file path: {}".format(rp))
        if re.match(".*/work/.{2}/.{30}/.*", rp):
            os.remove(rp)

@MainWrapper
def main(sample_id, fastq_pair, clear):
    """Main executor of the skesa template.

    Parameters
    ----------
    sample_id : str
        Sample Identification string.
    fastq_pair : list
        Two element list containing the paired FastQ files.
    clear : str
        Can be either 'true' or 'false'. If 'true', the input fastq files will
        be removed at the end of the run, IF they are in the working directory
    """

    logger.info("Starting skesa")

    # Determine output file
    if "_trim." in fastq_pair[0]:
        sample_id += "_trim"
    version = __get_version_skesa()["version"]
    output_file = "{}_skesa{}.fasta".format(sample_id, version.replace(".", ""))

    cli = [
        "skesa",
        "--fastq",
        "{},{}".format(fastq_pair[0], fastq_pair[1]),
        "--gz",
        "--use_paired_ends",
        "--cores",
        "${task.cpus}"
    ]

    logger.debug("Running Skesa subprocess with command: {}".format(cli))

    with open(output_file, "w") as fh:
        p = subprocess.Popen(cli, stdout=fh, stderr=PIPE)
    stdout, stderr = p.communicate()

    # Attempt to decode STDERR output from bytes. If unsuccessful, coerce to
    # string
    try:
        stderr = stderr.decode("utf8")
        stdout = stdout.decode("utf8")
    except (UnicodeDecodeError, AttributeError):
        stderr = str(stderr)
        stdout = str(stdout)

    logger.info("Finished Skesa subprocess with STDOUT:\\n"
                "======================================\\n{}".format(stdout))
    logger.info("Fished Skesa subprocess with STDERR:\\n"
                "======================================\\n{}".format(stderr))
    logger.info("Finished Skesa with return code: {}".format(
        p.returncode))

    # Remove input fastq files when clear option is specified.
    # Only remove temporary input when the expected output exists.
    if clear == "true" and os.path.exists(output_file):
        clean_up(fastq_pair)

    with open(".status", "w") as fh:
        if p.returncode != 0:
            fh.write("error")
            raise SystemExit(p.returncode)
        else:
            fh.write("pass")

if __name__ == '__main__':

    main(SAMPLE_ID, FASTQ_PAIR, CLEAR)


#!/usr/bin/env python3

"""
Purpose
-------

This module is intended parse the results of FastQC for paired end FastQ \
samples. It parses two reports:

    - Categorical report
    - Nucleotide level report.

Expected input
--------------

The following variables are expected whether using NextFlow or the
:py:func:`main` executor.

- ``sample_id`` : Sample identification string
    - e.g.: ``'SampleA'``

- ``result_p1`` : Path to both FastQC result files for pair 1
    - e.g.: ``'SampleA_1_data SampleA_1_summary'``

- ``result_p2`` : Path to both FastQC result files for pair 2
    - e.g.: ``'SampleA_2_data SampleA_2_summary'``

- ``opts`` : *Specify additional arguments for executing fastqc_report. \
    The arguments should be a string of command line arguments,\
    The accepted arguments are:*
    - ``'--ignore-tests'`` : Ignores test results from FastQC categorical\
    summary. This is used in the first run of FastQC.

Generated output
----------------

The generated output are output files that contain an object, usually a string.

- ``fastqc_health`` : Stores the health check for the current sample. If it
    passes all checks, it contains only the string 'pass'. Otherwise, contains
    the summary categories and their respective results
    - e.g.: ``'pass'``
- ``optimal_trim`` : Stores a tuple with the optimal trimming positions for 5'
    and 3' ends of the reads.
    - e.g.: ``'15 151'``

Code documentation
------------------

"""

__version__ = "1.0.2"
__build__ = "12052018"
__template__ = "fastqc_report-nf"

import os
import json

from collections import OrderedDict

from flowcraft_utils.flowcraft_base import get_logger, MainWrapper

logger = get_logger(__file__)

if __file__.endswith(".command.sh"):
    RESULT_P1 = '$result_p1'.split()
    RESULT_P2 = '$result_p2'.split()
    SAMPLE_ID = '$sample_id'
    OPTS = '$opts'.split()
    logger.debug("Running {} with parameters:".format(
        os.path.basename(__file__)))
    logger.debug("RESULT_P1: {}".format(RESULT_P1))
    logger.debug("RESULT_P2: {}".format(RESULT_P2))
    logger.debug("SAMPLE_ID: {}".format(SAMPLE_ID))
    logger.debug("OPTS: {}".format(OPTS))

def _get_quality_stats(d, start_str, field_start=1, field_end=2):
    """

    Parameters
    ----------
    d

    Returns
    -------

    """

    min_parsed = False
    parse = False
    report = []
    start_str = start_str
    end_str = ">>END_MODULE"

    with open(d) as fh:

        for line in fh:

            if line.startswith(start_str):
                next(fh)
                parse = True
                status = line.strip().split()[-1]

            # Exit parser when end string is found
            elif parse and line.startswith(end_str):
                return report, status

            elif parse:

                fields = line.strip().split()

                # This is triggered when the first value of a line series is
                # not 1. If the starting point of the series is a number
                # different from 1, fill the report with 0 until that point
                if not min_parsed:
                    if fields[0] != "1":
                        try:
                            blank_points = int(fields[0]) - 1
                            report.extend([0] * blank_points)
                        except ValueError:
                            pass
                    min_parsed = True

                report.append(";".join([
                    str(round(float(x), 2)) for x in
                    fields[field_start: field_end]
                ]))

def write_json_report(sample_id, data1, data2):
    """Writes the report

    Parameters
    ----------
    data1
    data2

    Returns
    -------

    """

    parser_map = {
        "base_sequence_quality": ">>Per base sequence quality",
        "sequence_quality": ">>Per sequence quality scores",
        "base_gc_content": ">>Per sequence GC content",
        "base_n_content": ">>Per base N content",
        "sequence_length_dist": ">>Sequence Length Distribution",
        "per_base_sequence_content": ">>Per base sequence content"
    }

    json_dic = {
        "plotData": [{
            "sample": sample_id,
            "data": {
                "base_sequence_quality": {"status": None, "data": []},
                "sequence_quality": {"status": None, "data": []},
                "base_gc_content": {"status": None, "data": []},
                "base_n_content": {"status": None, "data": []},
                "sequence_length_dist": {"status": None, "data": []},
                "per_base_sequence_content": {"status": None, "data": []}
            }
        }]
    }

    for cat, start_str in parser_map.items():

        if cat == "per_base_sequence_content":
            fs = 1
            fe = 5
        else:
            fs = 1
            fe = 2

        report1, status1 = _get_quality_stats(data1, start_str,
                                              field_start=fs, field_end=fe)
        report2, status2 = _get_quality_stats(data2, start_str,
                                              field_start=fs, field_end=fe)

        status = None
        for i in ["fail", "warn", "pass"]:
            if i in [status1, status2]:
                status = i

        json_dic["plotData"][0]["data"][cat]["data"] = [report1, report2]
        json_dic["plotData"][0]["data"][cat]["status"] = status

    return json_dic

def get_trim_index(biased_list):
    """Returns the trim index from a ``bool`` list

    Provided with a list of ``bool`` elements (``[False, False, True, True]``),
    this function will assess the index of the list that minimizes the number
    of True elements (biased positions) at the extremities. To do so,
    it will iterate over the boolean list and find an index position where
    there are two consecutive ``False`` elements after a ``True`` element. This
    will be considered as an optimal trim position. For example, in the
    following list::

        [True, True, False, True, True, False, False, False, False, ...]

    The optimal trim index will be the 4th position, since it is the first
    occurrence of a ``True`` element with two False elements after it.

    If the provided ``bool`` list has no ``True`` elements, then the 0 index is
    returned.

    Parameters
    ----------
    biased_list: list
        List of ``bool`` elements, where ``True`` means a biased site.

    Returns
    -------
        x : index position of the biased list for the optimal trim.

    """

    # Return index 0 if there are no biased positions
    if set(biased_list) == {False}:
        return 0

    if set(biased_list[:5]) == {False}:
        return 0

    # Iterate over the biased_list array. Keep the iteration going until
    # we find a biased position with the two following positions unbiased
    # (e.g.: True, False, False).
    # When this condition is verified, return the last biased position
    # index for subsequent trimming.
    for i, val in enumerate(biased_list):
        if val and set(biased_list[i+1:i+3]) == {False}:
            return i + 1

    # If the previous iteration could not find and index to trim, it means
    # that the whole list is basically biased. Return the length of the
    # biased_list
    return len(biased_list)

def trim_range(data_file):
    """Assess the optimal trim range for a given FastQC data file.

    This function will parse a single FastQC data file, namely the
    *'Per base sequence content'* category. It will retrieve the A/T and G/C
    content for each nucleotide position in the reads, and check whether the
    G/C and A/T proportions are between 80% and 120%. If they are, that
    nucleotide position is marked as biased for future removal.

    Parameters
    ----------
    data_file: str
        Path to FastQC data file.

    Returns
    -------
    trim_nt: list
        List containing the range with the best trimming positions for the
        corresponding FastQ file. The first element is the 5' end trim index
        and the second element is the 3' end trim index.
    """

    logger.debug("Starting trim range assessment")

    # Target string for nucleotide bias assessment
    target_nuc_bias = ">>Per base sequence content"
    logger.debug("Target string to start nucleotide bias assessment set to "
                 "{}".format(target_nuc_bias))
    # This flag will become True when gathering base proportion data
    # from file.
    gather = False

    # This variable will store a boolean array on the biased/unbiased
    # positions. Biased position will be True, while unbiased positions
    # will be False
    biased = []

    with open(data_file) as fh:

        for line in fh:
            # Start assessment of nucleotide bias
            if line.startswith(target_nuc_bias):
                # Skip comment line
                logger.debug("Found target string at line: {}".format(line))
                next(fh)
                gather = True
            # Stop assessment when reaching end of target module
            elif line.startswith(">>END_MODULE") and gather:
                logger.debug("Stopping parsing at line: {}".format(line))
                break
            elif gather:
                # Get proportions of each nucleotide
                g, a, t, c = [float(x) for x in line.strip().split()[1:]]
                # Get 'GC' and 'AT content
                gc = (g + 0.1) / (c + 0.1)
                at = (a + 0.1) / (t + 0.1)
                # Assess bias
                if 0.8 <= gc <= 1.2 and 0.8 <= at <= 1.2:
                    biased.append(False)
                else:
                    biased.append(True)

    logger.debug("Finished bias assessment with result: {}".format(biased))

    # Split biased list in half to get the 5' and 3' ends
    biased_5end, biased_3end = biased[:int(len(biased)/2)],\
        biased[int(len(biased)/2):][::-1]

    logger.debug("Getting optimal trim range from biased list")
    trim_nt = [0, 0]
    # Assess number of nucleotides to clip at 5' end
    trim_nt[0] = get_trim_index(biased_5end)
    logger.debug("Optimal trim range at 5' end set to: {}".format(trim_nt[0]))
    # Assess number of nucleotides to clip at 3' end
    trim_nt[1] = len(biased) - get_trim_index(biased_3end)
    logger.debug("Optimal trim range at 3' end set to: {}".format(trim_nt[1]))

    return trim_nt

def get_sample_trim(p1_data, p2_data):
    """Get the optimal read trim range from data files of paired FastQ reads.

    Given the FastQC data report files for paired-end FastQ reads, this
    function will assess the optimal trim range for the 3' and 5' ends of
    the paired-end reads. This assessment will be based on the *'Per sequence
    GC content'*.

    Parameters
    ----------
    p1_data: str
        Path to FastQC data report file from pair 1
    p2_data: str
        Path to FastQC data report file from pair 2

    Returns
    -------
    optimal_5trim: int
        Optimal trim index for the 5' end of the reads
    optima_3trim: int
        Optimal trim index for the 3' end of the reads

    See Also
    --------
    trim_range

    """

    sample_ranges = [trim_range(x) for x in [p1_data, p2_data]]

    # Get the optimal trim position for 5' end
    optimal_5trim = max([x[0] for x in sample_ranges])
    # Get optimal trim position for 3' end
    optimal_3trim = min([x[1] for x in sample_ranges])

    return optimal_5trim, optimal_3trim

def get_summary(summary_file):
    """Parses a FastQC summary report file and returns it as a dictionary.

    This function parses a typical FastQC summary report file, retrieving
    only the information on the first two columns. For instance, a line could
    be::

        'PASS	Basic Statistics	SH10762A_1.fastq.gz'

    This parser will build a dictionary with the string in the second column
    as a key and the QC result as the value. In this case, the returned
    ``dict`` would be something like::

        {"Basic Statistics": "PASS"}

    Parameters
    ----------
    summary_file: str
        Path to FastQC summary report.

    Returns
    -------
    summary_info: :py:data:`OrderedDict`
        Returns the information of the FastQC summary report as an ordered
        dictionary, with the categories as strings and the QC result as values.

    """

    summary_info = OrderedDict()
    logger.debug("Retrieving summary information from file: {}".format(
        summary_file))

    with open(summary_file) as fh:
        for line in fh:
            # Skip empty lines
            if not line.strip():
                continue
            # Populate summary info
            fields = [x.strip() for x in line.split("\t")]
            summary_info[fields[1]] = fields[0]

    logger.debug("Retrieved summary information from file: {}".format(
        summary_info))

    return summary_info

def check_summary_health(summary_file, **kwargs):
    """Checks the health of a sample from the FastQC summary file.

    Parses the FastQC summary file and tests whether the sample is good
    or not. There are four categories that cannot fail, and two that
    must pass in order for the sample pass this check. If the sample fails
    the quality checks, a list with the failing categories is also returned.

    Categories that cannot fail::

        fail_sensitive = [
            "Per base sequence quality",
            "Overrepresented sequences",
            "Sequence Length Distribution",
            "Per sequence GC content"
        ]

    Categories that must pass::

        must_pass = [
            "Per base N content",
            "Adapter Content"
        ]

    Parameters
    ----------
    summary_file: str
        Path to FastQC summary file.

    Returns
    -------
    x : bool
        Returns ``True`` if the sample passes all tests. ``False`` if not.
    summary_info : list
        A list with the FastQC categories that failed the tests. Is empty
        if the sample passes all tests.
    """

    # Store the summary categories that cannot fail. If they fail, do not
    # proceed with this sample
    fail_sensitive = kwargs.get("fail_sensitive", [
        "Per base sequence quality",
        "Overrepresented sequences",
        "Sequence Length Distribution",
        "Per sequence GC content"
    ])
    logger.debug("Fail sensitive categories: {}".format(fail_sensitive))

    # Store summary categories that must pass. If they do not, do not proceed
    # with that sample
    must_pass = kwargs.get("must_pass", [
        "Per base N content",
        "Adapter Content"
    ])
    logger.debug("Must pass categories: {}".format(must_pass))

    warning_fail_sensitive = kwargs.get("warning_fail_sensitive", [
        "Per base sequence quality",
        "Overrepresented sequences",

    ])

    warning_must_pass = kwargs.get("warning_must_pass", [
        "Per base sequence content"
    ])

    # Get summary dictionary
    summary_info = get_summary(summary_file)

    # This flag will change to False if one of the tests fails
    health = True
    # List of failing categories
    failed = []
    # List of warning categories
    warning = []

    for cat, test in summary_info.items():

        logger.debug("Assessing category {} with result {}".format(cat, test))

        # FAILURES
        # Check for fail sensitive
        if cat in fail_sensitive and test == "FAIL":
            health = False
            failed.append("{}:{}".format(cat, test))
            logger.error("Category {} failed a fail sensitive "
                         "category".format(cat))

        # Check for must pass
        if cat in must_pass and test != "PASS":
            health = False
            failed.append("{}:{}".format(cat, test))
            logger.error("Category {} failed a must pass category".format(
                cat))

        # WARNINGS
        # Check for fail sensitive
        if cat in warning_fail_sensitive and test == "FAIL":
            warning.append("Failed category: {}".format(cat))
            logger.warning("Category {} flagged at a fail sensitive "
                           "category".format(cat))

        if cat in warning_must_pass and test != "PASS":
            warning.append("Did not pass category: {}".format(cat))
            logger.warning("Category {} flagged at a must pass "
                           "category".format(cat))

    # Passed all tests
    return health, failed, warning

@MainWrapper
def main(sample_id, result_p1, result_p2, opts):
    """Main executor of the fastqc_report template.

    If the "--ignore-tests" option is present in the ``opts`` argument,
    the health check of the sample will be bypassed, and it will pass the
    check. This option is used in the first run of FastQC. In the second
    run (after filtering with trimmomatic) this option is not provided and
    the samples are submitted to a health check before proceeding in the
    pipeline.

    Parameters
    ----------
    sample_id : str
        Sample Identification string.
    result_p1 : list
        Two element list containing the path to the FastQC report files to
        the first FastQ pair.
        The first must be the nucleotide level report and the second the
        categorical report.
    result_p2: list
        Two element list containing the path to the FastQC report files to
        the second FastQ pair.
        The first must be the nucleotide level report and the second the
        categorical report.
    opts : list
        List of arbitrary options. See `Expected input`_.

    """

    logger.info("Starting fastqc report")
    json_dic = {}

    with open("{}_trim_report".format(sample_id), "w") as trep_fh, \
            open("optimal_trim", "w") as trim_fh, \
            open("{}_status_report".format(sample_id), "w") as rep_fh, \
            open(".status", "w") as status_fh, \
            open(".warning", "w") as warn_fh, \
            open(".fail", "w") as fail_fh, \
            open(".report.json", "w") as report_fh:

        # Perform health check according to the FastQC summary report for
        # each pair. If both pairs pass the check, send the 'pass' information
        # to the 'fastqc_health' channel. If at least one fails, send the
        # summary report.
        if "--ignore-tests" not in opts:

            # Get reports for each category in json format
            json_dic = write_json_report(sample_id, result_p1[0],
                                         result_p2[0])

            logger.info("Performing FastQ health check")
            for p, fastqc_summary in enumerate([result_p1[1], result_p2[1]]):

                logger.debug("Checking files: {}".format(fastqc_summary))
                # Get the boolean health variable and a list of failed
                # categories, if any
                health, f_cat, warnings = check_summary_health(fastqc_summary)
                logger.debug("Health checked: {}".format(health))
                logger.debug("Failed categories: {}".format(f_cat))

                # Write any warnings
                if warnings:
                    json_dic["warnings"] = [{
                        "sample": sample_id,
                        "table": "qc",
                        "value": []
                    }]
                    for w in warnings:
                        warn_fh.write("{}\\n".format(w))
                        json_dic["warnings"][0]["value"].append(w)

                # Rename category summary file to the channel that will publish
                # The results
                output_file = "{}_{}_summary.txt".format(sample_id, p)
                os.rename(fastqc_summary, output_file)
                logger.debug("Setting summary file name to {}".format(
                    output_file))

                # If one of the health flags returns False, send the summary
                # report through the status channel
                if not health:
                    fail_msg = "Sample failed quality control checks:" \
                               " {}".format(",".join(f_cat))
                    logger.warning(fail_msg)
                    fail_fh.write(fail_msg)
                    json_dic["fail"] = [{
                        "sample": sample_id,
                        "table": "qc",
                        "value": [fail_msg]
                    }]
                    report_fh.write(
                        json.dumps(json_dic, separators=(",", ":")))
                    status_fh.write("fail")
                    trim_fh.write("fail")
                    rep_fh.write("{}, {}\\n".format(sample_id, ",".join(f_cat)))
                    trep_fh.write("{},fail,fail\\n".format(sample_id))

                    return

            logger.info("Sample passed quality control checks")

        status_fh.write("pass")
        rep_fh.write("{}, pass\\n".format(sample_id))

        logger.info("Assessing optimal trim range for sample")
        # Get optimal trimming range for sample, based on the per base sequence
        # content
        optimal_trim = get_sample_trim(result_p1[0], result_p2[0])
        logger.info("Optimal trim range set to: {}".format(optimal_trim))
        trim_fh.write("{}".format(" ".join([str(x) for x in optimal_trim])))

        trep_fh.write("{},{},{}\\n".format(sample_id, optimal_trim[0],
                                           optimal_trim[1]))

        # The json dict report is only populated when the FastQC quality
        # checks are performed, that is, when the --ignore-tests option
        # is not provide
        if json_dic:
            report_fh.write(json.dumps(json_dic, separators=(",", ":")))

if __name__ == '__main__':

    main(SAMPLE_ID, RESULT_P1, RESULT_P2, OPTS)


#!/usr/bin/env python3

import re
import os
import json

from flowcraft_utils.flowcraft_base import get_logger, MainWrapper

"""
Purpose
-------

This module is intended to process the output of mapping proces from a single
sample from the program Bowtie for the report component.
The main input is an log file produced by the mapper.

Expected input
--------------

The following variables are expected whether using NextFlow or the
:py:func:`main` executor.

- ``sample_id``: Sample Identification string.
    - e.g.: ``'SampleA'``
- ``bowtie_log``: Log file from the mapper.
    - e.g.: ``'bowtie.log'``

Generated output
----------------
- ``.report.jason``: Data structure for the report

Code documentation
------------------

"""

__version__ = "1.0.1"
__build__ = "10.09.2018"
__template__ = "remove_host-nf"

logger = get_logger(__file__)

if __file__.endswith(".command.sh"):
    SAMPLE_ID = '$sample_id'
    BOWTIE_LOG = '$bowtie_log'
    logger.debug("Running {} with parameters:".format(
        os.path.basename(__file__)))
    logger.debug("SAMPLE_ID: {}".format(SAMPLE_ID))
    logger.debug("BOWTIE_LOG: {}".format(BOWTIE_LOG))

class Bowtie:
    """
    Class to parse and store the info in the bowtie log file.

    """

    def __init__(self, sample_id, bowtie_log):

        self.sample = sample_id
        """
        str: The name of the sample for the assembly.
        """

        self.n_reads = 0

        self.align_0x = 0

        self.align_1x = 0

        self.align_mt1x = 0

        self.overall_rate = 0.0

        # Parse assembly and populate self.n_reads, self.align_0x, self.align_1x, self.align_mt1x and self.overall_rate
        self.parse_log(bowtie_log)

    def set_n_reads(self, n_reads):
        self.n_reads = int(n_reads)

    def set_align_0x(self,align_0x):
        self.align_0x = align_0x

    def set_align_1x(self,align_1x):
        self.align_1x = align_1x

    def set_align_mt1x(self,align_mt1x):
        self.align_mt1x = align_mt1x

    def set_overall_rate(self,overall_rate):
        self.overall_rate = overall_rate

    def parse_log(self, bowtie_log):
        """Parse a bowtie log file.

        This is a bowtie log parsing method that populates the
        :py:attr:`self.n_reads, self.align_0x, self.align_1x, self.align_mt1x and self.overall_rate` attributes with
        data from the log file.

        Disclamer: THIS METHOD IS HORRIBLE BECAUSE THE BOWTIE LOG IS HORRIBLE.

        The insertion of data on the attribytes is done by the
        :py:meth:`set_attribute method.

        Parameters
        ----------
        bowtie_log : str
            Path to the boetie log file.

       """

        print("is here!")

        # Regexes - thanks to https://github.com/ewels/MultiQC/blob/master/multiqc/modules/bowtie2/bowtie2.py
        regexes = {
            'unpaired': {
                'unpaired_aligned_none': r"(\\d+) \\([\\d\\.]+%\\) aligned 0 times",
                'unpaired_aligned_one': r"(\\d+) \\([\\d\\.]+%\\) aligned exactly 1 time",
                'unpaired_aligned_multi': r"(\\d+) \\([\\d\\.]+%\\) aligned >1 times"
            },
            'paired': {
                'paired_aligned_none': r"(\\d+) \\([\\d\\.]+%\\) aligned concordantly 0 times",
                'paired_aligned_one': r"(\\d+) \\([\\d\\.]+%\\) aligned concordantly exactly 1 time",
                'paired_aligned_multi': r"(\\d+) \\([\\d\\.]+%\\) aligned concordantly >1 times",
                'paired_aligned_discord_one': r"(\\d+) \\([\\d\\.]+%\\) aligned discordantly 1 time",
                'paired_aligned_discord_multi': r"(\\d+) \\([\\d\\.]+%\\) aligned discordantly >1 times",
                'paired_aligned_mate_one': r"(\\d+) \\([\\d\\.]+%\\) aligned exactly 1 time",
                'paired_aligned_mate_multi': r"(\\d+) \\([\\d\\.]+%\\) aligned >1 times",
                'paired_aligned_mate_none': r"(\\d+) \\([\\d\\.]+%\\) aligned 0 times"
            }
        }

        #Missing parser for unpaired (not implemented in flowcraft yet)

        with open(bowtie_log, "r") as f:
            #Go through log file line by line
            for l in f:

                print(l)

                #total reads
                total = re.search(r"(\\d+) reads; of these:", l)
                print(total)
                if total:
                    print(total)
                    self.set_n_reads(total.group(1))

                # Paired end reads aka the pain
                paired = re.search(r"(\\d+) \\([\\d\\.]+%\\) were paired; of these:", l)
                if paired:
                    paired_total = int(paired.group(1))

                    paired_numbers = {}

                    # Do nested loop whilst we have this level of indentation
                    l = f.readline()
                    while l.startswith('    '):
                        for k, r in regexes['paired'].items():
                            match = re.search(r, l)
                            if match:
                                paired_numbers[k] = int(match.group(1))
                        l = f.readline()

                    align_zero_times = paired_numbers['paired_aligned_none'] + paired_numbers['paired_aligned_mate_none']
                    if align_zero_times:
                        self.set_align_0x(align_zero_times)

                    align_one_time = paired_numbers['paired_aligned_one'] + paired_numbers['paired_aligned_mate_one']
                    if align_one_time:
                        self.set_align_1x(align_one_time)

                    align_more_than_one_time = paired_numbers['paired_aligned_multi'] + paired_numbers['paired_aligned_mate_multi']
                    if align_more_than_one_time:
                        self.set_align_mt1x(align_more_than_one_time)

                # Overall alignment rate
                overall = re.search(r"([\\d\\.]+)% overall alignment rate", l)
                if overall:
                    self.overall_rate = float(overall.group(1))

@MainWrapper
def main(sample_id, bowite_log):
    """Main executor of the process_mapping template.

    Parameters
    ----------
    sample_id : str
        Sample Identification string.
    boetie_log: str
        Path to the log file generated by bowtie.

    """

    logger.info("Starting mapping file processing")
    warnings = []
    fails = ""

    bowtie_info = Bowtie(sample_id, bowite_log)

    print(bowtie_info.overall_rate)

    with open(".report.json", "w") as json_report:
        json_dic = {
            "tableRow": [{
                "sample": sample_id,
                "data": [
                    {"header": "Reads",
                     "value": int(bowtie_info.n_reads),
                     "table": "mapping",
                     "columnBar": False},
                    {"header": "Unmapped",
                     "value": int(bowtie_info.align_0x),
                     "table": "mapping",
                     "columnBar": False},
                    {"header": "Mapped 1x",
                     "value": int(bowtie_info.align_1x),
                     "table": "mapping",
                     "columnBar": False},
                    {"header": "Mapped >1x",
                     "value": int(bowtie_info.align_mt1x),
                     "table": "mapping",
                     "columnBar": False},
                    {"header": "Overall alignment rate (%)",
                     "value": float(bowtie_info.overall_rate),
                     "table": "mapping",
                     "columnBar": False}
                ]
            }],
        }

        if warnings:
            json_dic["warnings"] = [{
                "sample": sample_id,
                "table": "mapping",
                "value": warnings
            }]

        if fails:
            json_dic["fail"] = [{
                "sample": sample_id,
                "table": "mapping",
                "value": [fails]
            }]

        json_report.write(json.dumps(json_dic, separators=(",", ":")))

    with open(".status", "w") as status_fh:
        status_fh.write("pass")

if __name__ == '__main__':

    main(SAMPLE_ID, BOWTIE_LOG)