"""

"""

import os
import sys
import json
import logging
import traceback

from time import gmtime, strftime

def get_logger(filepath, level=logging.DEBUG):
    # create logger
    logger = logging.getLogger(os.path.basename(filepath))
    logger.setLevel(level)
    # create console handler and set level to debug
    ch = logging.StreamHandler()
    ch.setLevel(level)
    # create formatter
    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    # add formatter to ch
    ch.setFormatter(formatter)
    # add ch to logger
    logger.addHandler(ch)

    return logger

def log_error():
    """Nextflow specific function that logs an error upon unexpected failing
    """

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

class MainWrapper:

    def __init__(self, f):

        self.f = f
        self.context = self.f.__globals__
        self.logger = self.context.get("logger", None)

    def __call__(self, *args, **kwargs):

        self.logger.debug("Starting template at {}".format(
            strftime("%Y-%m-%d %H:%M:%S", gmtime())))
        self.logger.debug("Working directory: {}".format(os.getcwd()))

        try:
            self.build_versions()
            self.f(*args, **kwargs)
        except SystemExit as e:
            sys.exit(e)
        except:
            if self.logger:
                self.logger.error("Module exited unexpectedly with error:"
                                  "\\n{}".format(traceback.format_exc()))
            log_error()

        self.logger.debug("Finished template at {}".format(
            strftime("%Y-%m-%d %H:%M:%S", gmtime())))

    def build_versions(self):
        """Writes versions JSON for a template file

        This method creates the JSON file ``.versions`` based on the metadata
        and specific functions that are present in a given template script.

        It starts by fetching the template metadata, which can be specified
        via the ``__version__``, ``__template__`` and ``__build__``
        attributes. If all of these attributes exist, it starts to populate
        a JSON/dict array (Note that the absence of any one of them will
        prevent the version from being written).

        Then, it will search the
        template scope for functions that start with the substring
        ``__set_version`` (For example ``def __set_version_fastqc()`).
        These functions should gather the version of
        an arbitrary program and return a JSON/dict object with the following
        information::

            {
                "program": <program_name>,
                "version": <version>
                "build": <build>
            }

        This JSON/dict object is then written in the ``.versions`` file.
        """

        version_storage = []

        template_version = self.context.get("__version__", None)
        template_program = self.context.get("__template__", None)
        template_build = self.context.get("__build__", None)

        if template_version and template_program and template_build:
            if self.logger:
                self.logger.debug("Adding template version: {}; {}; "
                                  "{}".format(template_program,
                                              template_version,
                                              template_build))
            version_storage.append({
                "program": template_program,
                "version": template_version,
                "build": template_build
            })

        for var, obj in self.context.items():
            if var.startswith("__get_version"):
                ver = obj()
                version_storage.append(ver)
                if self.logger:
                    self.logger.debug("Found additional software version"
                                      "{}".format(ver))

        with open(".versions", "w") as fh:
            fh.write(json.dumps(version_storage, separators=(",", ":")))


#!/usr/bin/env python3

"""
Purpose
-------

This module is intended to generate a json output for mash screen results that
can be imported in pATLAS.

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

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

- ``mash_output`` : String with the name of the mash screen output file.
    - e.g.: ``'sortedMashScreenResults_SampleA.txt'``

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

"""

__version__ = "1.1.0"
__build__ = "04072018"
__template__ = "mashscreen2json-nf"

from statistics import median
import os
import json

from flowcraft_utils.flowcraft_base import get_logger, MainWrapper

logger = get_logger(__file__)

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

@MainWrapper
def main(mash_output, sample_id):
    '''
    converts top results from mash screen txt output to json format

    Parameters
    ----------
    mash_output: str
        this is a string that stores the path to this file, i.e, the name of
        the file
    sample_id: str
        sample name

    '''
    logger.info("Reading file : {}".format(mash_output))
    read_mash_output = open(mash_output)

    dic = {}
    median_list = []
    filtered_dic = {}

    logger.info("Generating dictionary and list to pre-process the final json")
    for line in read_mash_output:
        tab_split = line.split("\t")
        identity = tab_split[0]
        # shared_hashes = tab_split[1]
        median_multiplicity = tab_split[2]
        # p_value = tab_split[3]
        query_id = tab_split[4]
        # query-comment should not exist here and it is irrelevant

        # here identity is what in fact interests to report to json but
        # median_multiplicity also is important since it gives an rough
        # estimation of the coverage depth for each plasmid.
        # Plasmids should have higher coverage depth due to their increased
        # copy number in relation to the chromosome.
        dic[query_id] = [identity, median_multiplicity]
        median_list.append(float(median_multiplicity))

    output_json = open(" ".join(mash_output.split(".")[:-1]) + ".json", "w")

    # median cutoff is twice the median of all median_multiplicity values
    # reported by mash screen. In the case of plasmids, since the database
    # has 9k entries and reads shouldn't have that many sequences it seems ok...
    if len(median_list) > 0:
        # this statement assures that median_list has indeed any entries
        median_cutoff = median(median_list)
        logger.info("Generating final json to dump to a file")
        for k, v in dic.items():
            # estimated copy number
            copy_number = int(float(v[1]) / median_cutoff)
            # assure that plasmid as at least twice the median coverage depth
            if float(v[1]) > median_cutoff:
                filtered_dic["_".join(k.split("_")[0:3])] = [
                    round(float(v[0]),2),
                    copy_number
                ]
        logger.info(
            "Exported dictionary has {} entries".format(len(filtered_dic)))
    else:
        # if no entries were found raise an error
        logger.error("No matches were found using mash screen for the queried reads")

    output_json.write(json.dumps(filtered_dic))
    output_json.close()

    json_dic = {
        "tableRow": [{
            "sample": sample_id,
            "data": [{
                "header": "Mash Screen",
                "table": "plasmids",
                "patlas_mashscreen": filtered_dic,
                "value": len(filtered_dic)
            }]
        }],
    }

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

if __name__ == "__main__":

    main(MASH_TXT, SAMPLE_ID)


import logging
import sys

logger = logging.getLogger("main.{}".format(__name__))

COLORS = {
    "green_bold": "1;32m",
    "red_bold": "1;31m",
    "white": "0;38m",
    "white_bold": "1;38m",
    "white_underline": "4;38m",
    "blue_bold": "1;36m",
    "purple_bold": "1;34m",
    "yellow_bold": "1;93m"
}

def colored_print(msg, color_label="white_bold"):
    """
    This function enables users to add a color to the print. It also enables
    to pass end_char to print allowing to print several strings in the same line
    in different prints.

    Parameters
    ----------
    color_string: str
        The color code to pass to the function, which enables color change as
        well as background color change.
    msg: str
        The actual text to be printed
    end_char: str
        The character in which each print should finish. By default it will be
        "\n".

    """

    if sys.stdout.encoding != "UTF-8":
        msg = "".join([i if ord(i) < 128 else "" for i in msg])

    # try except first looks for the color in COLORS dictionary, otherwise use
    # color_label as the color.
    try:
        col = COLORS[color_label]
    except KeyError:
        col = color_label

    return "\x1b[{}{}\x1b[0m".format(col, msg)

def procs_dict_parser(procs_dict):
    """
    This function handles the dictionary of attributes of each Process class
    to print to stdout lists of all the components or the components which the
    user specifies in the -t flag.

    Parameters
    ----------
    procs_dict: dict
        A dictionary with the class attributes for all the components (or
        components that are used by the -t flag), that allow to create
        both the short_list and detailed_list. Dictionary example:
        {"abyss": {'input_type': 'fastq', 'output_type': 'fasta',
        'dependencies': [], 'directives': {'abyss': {'cpus': 4,
        'memory': '{ 5.GB * task.attempt }', 'container': 'flowcraft/abyss',
        'version': '2.1.1', 'scratch': 'true'}}}
    """

    logger.info(colored_print(
        "\n===== L I S T   O F   P R O C E S S E S =====\n", "green_bold"))

    #Sort to print alphabetically ordered list of processes to ease reading
    procs_dict_ordered = {k: procs_dict[k] for k in sorted(procs_dict)}

    for template, dict_proc_info in procs_dict_ordered.items():
        template_str = "=> {}".format(template)
        logger.info(colored_print(template_str, "blue_bold"))

        for info in dict_proc_info:
            info_str = "{}:".format(info)

            if isinstance(dict_proc_info[info], list):
                if not dict_proc_info[info]:
                    arg_msg = "None"
                else:
                    arg_msg = ", ".join(dict_proc_info[info])
            elif info == "directives":
                # this is used for the "directives", which is a dict
                if not dict_proc_info[info]:
                    # if dict is empty then add None to the message
                    arg_msg = "None"
                else:
                    # otherwise fetch all template names within a component
                    # and all the directives for each template to a list
                    list_msg = ["\n      {}: {}".format(
                        templt,
                        " , ".join(["{}: {}".format(dr, val)
                                    for dr, val in drs.items()]))
                                for templt, drs in dict_proc_info[info].items()
                    ]
                    # write list to a str
                    arg_msg = "".join(list_msg)
            else:
                arg_msg = dict_proc_info[info]

            logger.info("   {} {}".format(
                colored_print(info_str, "white_underline"), arg_msg
            ))

def proc_collector(process_map, args, pipeline_string):
    """
    Function that collects all processes available and stores a dictionary of
    the required arguments of each process class to be passed to
    procs_dict_parser

    Parameters
    ----------
    process_map: dict
        The dictionary with the Processes currently available in flowcraft
        and their corresponding classes as values
    args: argparse.Namespace
        The arguments passed through argparser that will be access to check the
        type of list to be printed
    pipeline_string: str
        the pipeline string

    """

    arguments_list = []

    # prints a detailed list of the process class arguments
    if args.detailed_list:
        # list of attributes to be passed to proc_collector
        arguments_list += [
            "input_type",
            "output_type",
            "description",
            "dependencies",
            "conflicts",
            "directives"
        ]

    # prints a short list with each process and the corresponding description
    if args.short_list:
        arguments_list += [
            "description"
        ]

    if arguments_list:
        # dict to store only the required entries
        procs_dict = {}
        # loops between all process_map Processes
        for name, cls in process_map.items():

            # instantiates each Process class
            cls_inst = cls(template=name)

            # checks if recipe is provided
            if pipeline_string:
                if name not in pipeline_string:
                    continue

            d = {arg_key: vars(cls_inst)[arg_key] for arg_key in
                 vars(cls_inst) if arg_key in arguments_list}
            procs_dict[name] = d

        procs_dict_parser(procs_dict)

        sys.exit(0)


#!/usr/bin/env python3

"""
Purpose
-------

This module receives paired FastQ files, a genome size estimate and a minimum
coverage threshold and has three purposes while iterating over the FastQ files:

    - Checks the integrity of FastQ files (corrupted files).
    - Guesses the encoding of FastQ files (this can be turned off in the \
    ``opts`` argument).
    - Estimates the coverage for each sample.

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'``

- ``gsize`` : *Expected genome size*
    - e.g.: ``'2.5'``

- ``cov`` : *Minimum coverage threshold*
    - e.g.: ``'15'``

- ``opts`` : *Specify additional arguments for executing integrity_coverage. \
    The arguments should be a string of command line arguments, such as \
    '-e'. The accepted arguments are:*
    - ``'-e'`` : Skip encoding guess.

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

The generated output are output files that contain an object, usually a string.
(Values within ``${}`` are substituted by the corresponding variable.)

- ``${sample_id}_encoding`` : Stores the encoding for the sample FastQ. If no \
    encoding could be guessed, write 'None' to file.
    - e.g.: ``'Illumina-1.8'`` or ``'None'``

- ``${sample_id}_phred`` : Stores the phred value for the sample FastQ. If no \
    phred could be guessed, write 'None' to file.
    - ``'33'`` or ``'None'``

- ``${sample_id}_coverage`` : Stores the expected coverage of the samples, \
    based on a given genome size.
    - ``'112'`` or ``'fail'``

- ``${sample_id}_report`` : Stores the report on the expected coverage \
    estimation. This string written in this file will appear in the \
    coverage report.
    - ``'${sample_id}, 112, PASS'``

- ``${sample_id}_max_len`` : Stores the maximum read length for the current \
    sample.
    - ``'152'``

Notes
-----

In case of a corrupted sample, all expected output files should have
``'corrupt'`` written.

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

"""

__version__ = "1.0.1"
__build__ = "03082018"
__template__ = "integrity_coverage-nf"

import os
import bz2
import gzip
import json
import zipfile

from itertools import chain

from flowcraft_utils.flowcraft_base import get_logger, MainWrapper

logger = get_logger(__file__)

# Set constants when running from Nextflow
if __file__.endswith(".command.sh"):
    # CONSTANTS
    FASTQ_PAIR = '$fastq_pair'.split()
    SAMPLE_ID = '$sample_id'
    GSIZE = float('$gsize')
    MINIMUM_COVERAGE = float('$cov')
    OPTS = '$opts'.split()

    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("GSIZE: {}".format(GSIZE))
    logger.debug("MINIMUM_COVERAGE: {}".format(MINIMUM_COVERAGE))
    logger.debug("OPTS: {}".format(OPTS))

RANGES = {
    'Sanger': [33, (33, 73)],
    'Illumina-1.8': [33, (33, 74)],
    'Solexa': [64, (59, 104)],
    'Illumina-1.3': [64, (64, 104)],
    'Illumina-1.5': [64, (66, 105)]
}
"""
dict: Dictionary containing the encoding values for several fastq formats. The
key contains the format and the value contains a list with the corresponding
phred score and a list with the range of encodings.
"""

COPEN = {
    "gz": gzip.open,
    "bz2": bz2.open,
    "zip": zipfile.ZipFile
}

MAGIC_DICT = {
    b"\\x1f\\x8b\\x08": "gz",
    b"\\x42\\x5a\\x68": "bz2",
    b"\\x50\\x4b\\x03\\x04": "zip"
}
"""
dict: Dictionary containing the binary signatures for three compression formats
(gzip, bzip2 and zip).
"""

def guess_file_compression(file_path, magic_dict=None):
    """Guesses the compression of an input file.

    This function guesses the compression of a given file by checking for
    a binary signature at the beginning of the file. These signatures are
    stored in the :py:data:`MAGIC_DICT` dictionary. The supported compression
    formats are gzip, bzip2 and zip. If none of the signatures in this
    dictionary are found at the beginning of the file, it returns ``None``.

    Parameters
    ----------
    file_path : str
        Path to input file.
    magic_dict : dict, optional
        Dictionary containing the signatures of the compression types. The
        key should be the binary signature and the value should be the
        compression format. If left ``None``, it falls back to
        :py:data:`MAGIC_DICT`.

    Returns
    -------
    file_type : str or None
        If a compression type is detected, returns a string with the format.
        If not, returns ``None``.
    """

    if not magic_dict:
        magic_dict = MAGIC_DICT

    max_len = max(len(x) for x in magic_dict)

    with open(file_path, "rb") as f:
        file_start = f.read(max_len)

    logger.debug("Binary signature start: {}".format(file_start))

    for magic, file_type in magic_dict.items():
        if file_start.startswith(magic):
            return file_type

    return None

def get_qual_range(qual_str):
    """ Get range of the Unicode encode range for a given string of characters.

    The encoding is determined from the result of the :py:func:`ord` built-in.

    Parameters
    ----------
    qual_str : str
        Arbitrary string.

    Returns
    -------
    x : tuple
        (Minimum Unicode code, Maximum Unicode code).
    """

    vals = [ord(c) for c in qual_str]

    return min(vals), max(vals)

def get_encodings_in_range(rmin, rmax):
    """ Returns the valid encodings for a given encoding range.

    The encoding ranges are stored in the :py:data:`RANGES` dictionary, with
    the encoding name as a string and a list as a value containing the
    phred score and a tuple with the encoding range. For a given encoding
    range provided via the two first arguments, this function will return
    all possible encodings and phred scores.

    Parameters
    ----------
    rmin : int
        Minimum Unicode code in range.
    rmax : int
        Maximum Unicode code in range.

    Returns
    -------
    valid_encodings : list
        List of all possible encodings for the provided range.
    valid_phred : list
        List of all possible phred scores.

    """

    valid_encodings = []
    valid_phred = []

    for encoding, (phred, (emin, emax)) in RANGES.items():
        if rmin >= emin and rmax <= emax:
            valid_encodings.append(encoding)
            valid_phred.append(phred)

    return valid_encodings, valid_phred

@MainWrapper
def main(sample_id, fastq_pair, gsize, minimum_coverage, opts):
    """ Main executor of the integrity_coverage template.

    Parameters
    ----------
    sample_id : str
        Sample Identification string.
    fastq_pair : list
        Two element list containing the paired FastQ files.
    gsize : float or int
        Estimate of genome size in Mb.
    minimum_coverage : float or int
        Minimum coverage required for a sample to pass the coverage check
    opts : list
        List of arbitrary options. See `Expected input`_.

    """

    logger.info("Starting integrity coverage main")

    # Check for runtime options
    if "-e" in opts:
        skip_encoding = True
    else:
        skip_encoding = False

    # Information for encoding guess
    gmin, gmax = 99, 0
    encoding = []
    phred = None

    # Information for coverage estimation
    chars = 0
    nreads = 0

    # Information on maximum read length
    max_read_length = 0

    # Get compression of each FastQ pair file
    file_objects = []
    for fastq in fastq_pair:

        logger.info("Processing file {}".format(fastq))

        logger.info("[{}] Guessing file compression".format(fastq))
        ftype = guess_file_compression(fastq)

        # This can guess the compression of gz, bz2 and zip. If it cannot
        # find the compression type, it tries to open a regular file
        if ftype:
            logger.info("[{}] Found file compression: {}".format(
                fastq, ftype))
            file_objects.append(COPEN[ftype](fastq, "rt"))
        else:
            logger.info("[{}] File compression not found. Assuming an "
                        "uncompressed file".format(fastq))
            file_objects.append(open(fastq))

    logger.info("Starting FastQ file parsing")

    # The '*_encoding' file stores a string with the encoding ('Sanger')
    # If no encoding is guessed, 'None' should be stored
    # The '*_phred' file stores a string with the phred score ('33')
    # If no phred is guessed, 'None' should be stored
    # The '*_coverage' file stores the estimated coverage ('88')
    # The '*_report' file stores a csv report of the file
    # The '*_max_len' file stores a string with the maximum contig len ('155')
    with open("{}_encoding".format(sample_id), "w") as enc_fh, \
            open("{}_phred".format(sample_id), "w") as phred_fh, \
            open("{}_coverage".format(sample_id), "w") as cov_fh, \
            open("{}_report".format(sample_id), "w") as cov_rep, \
            open("{}_max_len".format(sample_id), "w") as len_fh, \
            open(".report.json", "w") as json_report, \
            open(".status", "w") as status_fh, \
            open(".fail", "w") as fail_fh:

        try:
            # Iterate over both pair files sequentially using itertools.chain
            for i, line in enumerate(chain(*file_objects)):

                # Parse only every 4th line of the file for the encoding
                # e.g.: AAAA/EEEEEEEEEEE<EEEEEEEEEEEEEEEEEEEEEEEEE (...)
                if (i + 1) % 4 == 0 and not skip_encoding:
                    # It is important to strip() the line so that any newline
                    # character is removed and not accounted for in the
                    # encoding guess
                    lmin, lmax = get_qual_range(line.strip())

                    # Guess new encoding if the range expands the previously
                    # set boundaries of gmin and gmax
                    if lmin < gmin or lmax > gmax:
                        gmin, gmax = min(lmin, gmin), max(lmax, gmax)
                        encoding, phred = get_encodings_in_range(gmin, gmax)
                        logger.debug(
                            "Updating estimates at line {} with range {} to"
                            " '{}' (encoding) and '{}' (phred)".format(
                                i, [lmin, lmax], encoding, phred))

                # Parse only every 2nd line of the file for the coverage
                # e.g.: GGATAATCTACCTTGACGATTTGTACTGGCGTTGGTTTCTTA (...)
                if (i + 3) % 4 == 0:
                    read_len = len(line.strip())
                    chars += read_len
                    nreads += 1

                    # Evaluate maximum read length for sample
                    if read_len > max_read_length:
                        logger.debug("Updating maximum read length at line "
                                     "{} to {}".format(i, read_len))
                        max_read_length = read_len

            # End of FastQ parsing
            logger.info("Finished FastQ file parsing")

            # The minimum expected coverage for a sample to pass
            exp_coverage = round(chars / (gsize * 1e6), 2)

            # Set json report
            if "-e" not in opts:

                json_dic = {
                    "tableRow": [{
                        "sample": sample_id,
                        "data": [
                            {"header": "Raw BP",
                             "value": chars,
                             "table": "qc",
                             "columnBar": True},
                            {"header": "Reads",
                             "value": nreads,
                             "table": "qc",
                             "columnBar": True},
                            {"header": "Coverage",
                             "value": exp_coverage,
                             "table": "qc",
                             "columnBar": True,
                             "failThreshold": minimum_coverage
                             }
                        ]
                    }],
                    "plotData": [{
                        "sample": sample_id,
                        "data": {
                            "sparkline": chars
                        }
                    }],
                }
            else:
                json_dic = {
                    "tableRow": [{
                        "sample": sample_id,
                        "data": [
                            {"header": "Coverage",
                             "value": exp_coverage,
                             "table": "qc",
                             "columnBar": True,
                             "failThreshold": minimum_coverage
                             }
                        ],
                    }],
                }

            # Get encoding
            if len(encoding) > 0:
                encoding = set(encoding)
                phred = set(phred)
                # Get encoding and phred as strings
                # e.g. enc: Sanger, Illumina-1.8
                # e.g. phred: 64
                enc = "{}".format(",".join([x for x in encoding]))
                phred = "{}".format(",".join(str(x) for x in phred))
                logger.info("Encoding set to {}".format(enc))
                logger.info("Phred set to {}".format(enc))

                enc_fh.write(enc)
                phred_fh.write(phred)
            # Encoding not found
            else:
                if not skip_encoding:
                    encoding_msg = "Could not guess encoding and phred from " \
                                   "FastQ"
                    logger.warning(encoding_msg)
                    json_dic["warnings"] = [{
                        "sample": sample_id,
                        "table": "qc",
                        "value": [encoding_msg]
                    }]
                    enc_fh.write("None")
                    phred_fh.write("None")

            # Estimate coverage
            logger.info("Estimating coverage based on a genome size of "
                        "{}".format(gsize))
            logger.info("Expected coverage is {}".format(exp_coverage))

            if exp_coverage >= minimum_coverage:
                cov_rep.write("{},{},{}\\n".format(
                    sample_id, str(exp_coverage), "PASS"))
                cov_fh.write(str(exp_coverage))
                status_fh.write("pass")
            # Estimated coverage does not pass minimum threshold
            else:
                fail_msg = "Sample with low coverage ({}), below the {} " \
                           "threshold".format(exp_coverage, minimum_coverage)
                logger.error(fail_msg)
                fail_fh.write(fail_msg)
                cov_fh.write("fail")
                status_fh.write("fail")
                cov_rep.write("{},{},{}\\n".format(
                    sample_id, str(exp_coverage), "FAIL"))
                json_dic["fail"] = [{
                    "sample": sample_id,
                    "table": "qc",
                    "value": [fail_msg]
                }]

            json_report.write(json.dumps(json_dic, separators=(",", ":")))
            # Maximum read length
            len_fh.write("{}".format(max_read_length))

        # This exception is raised when the input FastQ files are corrupted
        except EOFError:
            logger.error("The FastQ files could not be correctly "
                         "parsed. They may be corrupt")
            for fh in [enc_fh, phred_fh, cov_fh, cov_rep, len_fh]:
                fh.write("corrupt")
                status_fh.write("fail")
                fail_fh.write("Could not read/parse FastQ. "
                              "Possibly corrupt file")

if __name__ == "__main__":

    main(SAMPLE_ID, FASTQ_PAIR, GSIZE, MINIMUM_COVERAGE, OPTS)


#!/usr/bin/env python3

"""
Purpose
-------

This module is intended to process the coverage report from the
:py:class:`assembly_mapping` process.

TODO: Better purpose

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`` : Fasta assembly file.
    - e.g.: ``'SH10761A.assembly.fasta'``
- ``coverage`` : TSV file with the average coverage for each assembled contig.
    - e.g.: ``'coverage.tsv'``
- ``coverage_bp`` : TSV file with the coverage for each assembled bp.
    - e.g.: ``'coverage.tsv'``
- ``bam_file`` : BAM file with the alignment of reads to the genome.
    - e.g.: ``'sorted.bam'``
- ``opts`` : List of options for processing assembly mapping output.
    1. Minimum coverage for assembled contigs. Can be``auto``.
        - e.g.: ``'auto'`` or ``'10'``
    2. Maximum number of contigs.
        - e.g.: '100'
- ``gsize``: Expected genome size.
    - e.g.: ``'2.5'``

Generated output
----------------
- ``${sample_id}_filtered.assembly.fasta`` : Filtered assembly file in Fasta \
    format.
    - e.g.: ``'SampleA_filtered.assembly.fasta'``
- ``filtered.bam`` : BAM file with the same filtering as the assembly file.
    - e.g.: ``filtered.bam``

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

"""

__version__ = "1.0.1"
__build__ = "09022018"
__template__ = "process_assembly_mapping-nf"

import os
import json
import shutil
import subprocess

from subprocess import PIPE
from collections import OrderedDict

from flowcraft_utils.flowcraft_base import get_logger, MainWrapper

logger = get_logger(__file__)

def __get_version_samtools():

    try:
        cli = ["samtools", "--version"]
        p = subprocess.Popen(cli, stdout=PIPE, stderr=PIPE)
        stdout = p.communicate()[0]

        version = stdout.splitlines()[0].split()[1].decode("utf8")
    except Exception as e:
        logger.debug(e)
        version = "undefined"

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

def __get_version_bowtie2():

    try:
        cli = ["bowtie2", "--version"]
        p = subprocess.Popen(cli, stdout=PIPE, stderr=PIPE)
        stdout = p.communicate()[0]

        version = stdout.splitlines()[0].split()[-1].decode("utf8")
    except Exception as e:
        logger.debug(e)
        version = "undefined"

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

if __file__.endswith(".command.sh"):
    SAMPLE_ID = '$sample_id'
    ASSEMBLY_FILE = '$assembly'
    COVERAGE_FILE = '$coverage'
    COVERAGE_BP_FILE = '$coverage_bp'
    BAM_FILE = '$bam_file'
    OPTS = [x.strip() for x in '$opts'.strip("[]").split(",")]
    GSIZE = float('$gsize')
    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_FILE: {}".format(COVERAGE_FILE))
    logger.debug("COVERAGE_BP_FILE: {}".format(COVERAGE_BP_FILE))
    logger.debug("BAM_FILE: {}".format(BAM_FILE))
    logger.debug("MIN_ASSEMBLY_COVERAGE: {}".format(OPTS))
    logger.debug("GSIZE: {}".format(GSIZE))

def parse_coverage_table(coverage_file):
    """Parses a file with coverage information into objects.

    This function parses a TSV file containing coverage results for
    all contigs in a given assembly and will build an ``OrderedDict``
    with the information about their coverage and length.  The length
    information is actually gathered from the contig header using a
    regular expression that assumes the usual header produced by Spades::

        contig_len = int(re.search("length_(.+?)_", line).group(1))

    Parameters
    ----------
    coverage_file : str
        Path to TSV file containing the coverage results.

    Returns
    -------
    coverage_dict : OrderedDict
        Contains the coverage and length information for each contig.
    total_size : int
        Total size of the assembly in base pairs.
    total_cov : int
        Sum of coverage values across all contigs.
    """

    # Stores the correspondence between a contig and the corresponding coverage
    # e.g.: {"contig_1": {"cov": 424} }
    coverage_dict = OrderedDict()
    # Stores the total coverage
    total_cov = 0

    with open(coverage_file) as fh:
        for line in fh:
            # Get contig and coverage
            contig, cov = line.strip().split()
            coverage_dict[contig] = {"cov": int(cov)}
            # Add total coverage
            total_cov += int(cov)
            logger.debug("Processing contig '{}' with coverage '{}'"
                         "".format(contig, cov))

    return coverage_dict, total_cov

def filter_assembly(assembly_file, minimum_coverage, coverage_info,
                    output_file):
    """Generates a filtered assembly file.

    This function generates a filtered assembly file based on an original
    assembly and a minimum coverage threshold.

    Parameters
    ----------
    assembly_file : str
        Path to original assembly file.
    minimum_coverage : int or float
        Minimum coverage required for a contig to pass the filter.
    coverage_info : OrderedDict or dict
        Dictionary containing the coverage information for each contig.
    output_file : str
        Path where the filtered assembly file will be generated.

    """

    # This flag will determine whether sequence data should be written or
    # ignored because the current contig did not pass the minimum
    # coverage threshold
    write_flag = False

    with open(assembly_file) as fh, open(output_file, "w") as out_fh:

        for line in fh:
            if line.startswith(">"):
                # Reset write_flag
                write_flag = False
                # Get header of contig
                header = line.strip()[1:]
                # Check coverage for current contig
                contig_cov = coverage_info[header]["cov"]
                # If the contig coverage is above the threshold, write to
                # output filtered assembly
                if contig_cov >= minimum_coverage:
                    write_flag = True
                    out_fh.write(line)

            elif write_flag:
                out_fh.write(line)

def filter_bam(coverage_info, bam_file, min_coverage, output_bam):
    """Uses Samtools to filter a BAM file according to minimum coverage

    Provided with a minimum coverage value, this function will use Samtools
    to filter a BAM file. This is performed to apply the same filter to
    the BAM file as the one applied to the assembly file in
    :py:func:`filter_assembly`.

    Parameters
    ----------
    coverage_info : OrderedDict or dict
        Dictionary containing the coverage information for each contig.
    bam_file : str
        Path to the BAM file.
    min_coverage : int
        Minimum coverage required for a contig to pass the filter.
    output_bam : str
        Path to the generated filtered BAM file.
    """

    # Get list of contigs that will be kept
    contig_list = [x for x, vals in coverage_info.items()
                   if vals["cov"] >= min_coverage]

    cli = [
        "samtools",
        "view",
        "-bh",
        "-F",
        "4",
        "-o",
        output_bam,
        "-@",
        "1",
        bam_file,
    ]

    cli += contig_list

    logger.debug("Runnig samtools view 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 samtools view subprocess with STDOUT:\\n"
                "======================================\\n{}".format(stdout))
    logger.info("Fished samtools view subprocesswith STDERR:\\n"
                "======================================\\n{}".format(stderr))
    logger.info("Finished samtools view with return code: {}".format(
        p.returncode))

    if not p.returncode:
        # Create index
        cli = [
            "samtools",
            "index",
            output_bam
        ]

        logger.debug("Runnig samtools index subprocess with command: "
                     "{}".format(cli))

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

        try:
            stderr = stderr.decode("utf8")
            stdout = stdout.decode("utf8")
        except (UnicodeDecodeError, AttributeError):
            stderr = str(stderr)
            stdout = str(stdout)

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

def check_filtered_assembly(coverage_info, coverage_bp, minimum_coverage,
                            genome_size, contig_size, max_contigs,
                            sample_id):
    """Checks whether a filtered assembly passes a size threshold

    Given a minimum coverage threshold, this function evaluates whether an
    assembly will pass the minimum threshold of ``genome_size * 1e6 * 0.8``,
    which means 80% of the expected genome size or the maximum threshold
    of ``genome_size * 1e6 * 1.5``, which means 150% of the expected genome
    size. It will issue a warning if any of these thresholds is crossed.
    In the case of an expected genome size below 80% it will return False.

    Parameters
    ----------
    coverage_info : OrderedDict or dict
        Dictionary containing the coverage information for each contig.
    coverage_bp : dict
        Dictionary containing the per base coverage information for each
        contig. Used to determine the total number of base pairs in the
        final assembly.
    minimum_coverage : int
        Minimum coverage required for a contig to pass the filter.
    genome_size : int
        Expected genome size.
    contig_size : dict
        Dictionary with the len of each contig. Contig headers as keys and
        the corresponding lenght as values.
    max_contigs : int
        Maximum threshold for contig number. A warning is issued if this
        threshold is crossed.
    sample_id : str
        Id or name of the current sample

    Returns
    -------
    x : bool
        True if the filtered assembly size is higher than 80% of the
        expected genome size.

    """

    # Get size of assembly after filtering contigs below minimum_coverage
    assembly_len = sum([v for k, v in contig_size.items()
                        if coverage_info[k]["cov"] >= minimum_coverage])
    logger.debug("Assembly length after filtering for minimum coverage of"
                 " {}: {}".format(minimum_coverage, assembly_len))
    # Get number of contigs after filtering
    ncontigs = len([x for x in coverage_info.values()
                    if x["cov"] >= minimum_coverage])
    logger.debug("Number of contigs: {}".format(ncontigs))
    # Get number of bp after filtering
    filtered_contigs = [k for k, v in coverage_info.items()
                        if v["cov"] >= minimum_coverage]
    logger.debug("Filtered contigs for minimum coverage of "
                 "{}: {}".format(minimum_coverage, filtered_contigs))
    total_assembled_bp = sum([sum(coverage_bp[x]) for x in filtered_contigs
                              if x in coverage_bp])
    logger.debug("Total number of assembled base pairs:"
                 "{}".format(total_assembled_bp))

    warnings = []
    fails = []
    health = True

    with open(".warnings", "w") as warn_fh, \
            open(".report.json", "w") as json_report:

        logger.debug("Checking assembly size after filtering : {}".format(
            assembly_len))

        # If the filtered assembly size is above the 150% genome size
        # threshold, issue a warning
        if assembly_len > genome_size * 1e6 * 1.5:
            warn_msg = "Assembly size ({}) smaller than the maximum" \
                       " threshold of 150% of expected genome size.".format(
                            assembly_len)
            logger.warning(warn_msg)
            warn_fh.write(warn_msg)
            fails.append("Large_genome_size_({})".format(assembly_len))

        # If the number of contigs in the filtered assembly size crosses the
        # max_contigs threshold, issue a warning
        logger.debug("Checking number of contigs: {}".format(
                len(coverage_info)))
        contig_threshold = max_contigs * genome_size / 1.5
        if ncontigs > contig_threshold:
            warn_msg = "The number of contigs ({}) exceeds the threshold of " \
                       "100 contigs per 1.5Mb ({})".format(
                            ncontigs, round(contig_threshold, 1))
            logger.warning(warn_msg)
            warn_fh.write(warn_msg)
            warnings.append(warn_msg)

        # If the filtered assembly size falls below the 80% genome size
        # threshold, fail this check and return False
        if assembly_len < genome_size * 1e6 * 0.8:
            warn_msg = "Assembly size smaller than the minimum" \
                       " threshold of 80% of expected genome size: {}".format(
                            assembly_len)
            logger.warning(warn_msg)
            warn_fh.write(warn_msg)
            fails.append("Small_genome_size_({})".format(assembly_len))
            assembly_len = sum([v for v in contig_size.values()])
            total_assembled_bp = sum(
                [sum(coverage_bp[x]) for x in coverage_info if x in
                 coverage_bp])
            logger.debug("Assembly length without coverage filtering: "
                         "{}".format(assembly_len))
            logger.debug("Total number of assembled base pairs without"
                         " filtering: {}".format(total_assembled_bp))

            health = False

        json_dic = {
            "plotData": [{
                "sample": sample_id,
                "data": {
                    "sparkline": total_assembled_bp
                }
            }]
        }

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

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

    return health

def get_coverage_from_file(coverage_file):
    """

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

    Returns
    -------

    """

    contig_coverage = {}

    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 contig_coverage:
                contig_coverage[header] = [coverage]
            else:
                contig_coverage[header].append(coverage)

    return contig_coverage

def evaluate_min_coverage(coverage_opt, assembly_coverage, assembly_size):
    """ Evaluates the minimum coverage threshold from the value provided in
    the coverage_opt.

    Parameters
    ----------
    coverage_opt : str or int or float
        If set to "auto" it will try to automatically determine the coverage
        to 1/3 of the assembly size, to a minimum value of 10. If it set
        to a int or float, the specified value will be used.
    assembly_coverage : int or float
        The average assembly coverage for a genome assembly. This value
        is retrieved by the `:py:func:parse_coverage_table` function.
    assembly_size : int
        The size of the genome assembly. This value is retrieved by the
        `py:func:get_assembly_size` function.

    Returns
    -------
    x: int
        Minimum coverage threshold.

    """

    if coverage_opt == "auto":
        # Get the 1/3 value of the current assembly coverage
        min_coverage = (assembly_coverage / assembly_size) * .3
        logger.info("Minimum assembly coverage automatically set to: "
                    "{}".format(min_coverage))
        # If the 1/3 coverage is lower than 10, change it to the minimum of
        # 10
        if min_coverage < 10:
            logger.info("Minimum assembly coverage cannot be set to lower"
                        " that 10. Setting to 10")
            min_coverage = 10
    else:
        min_coverage = int(coverage_opt)
        logger.info("Minimum assembly coverage manually set to: {}".format(
            min_coverage))

    return min_coverage

def get_assembly_size(assembly_file):
    """Returns the number of nucleotides and the size per contig for the
    provided assembly file path

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

    Returns
    -------
    assembly_size : int
        Size of the assembly in nucleotides
    contig_size : dict
        Length of each contig (contig name as key and length as value)

    """

    assembly_size = 0
    contig_size = {}
    header = ""

    with open(assembly_file) as fh:
        for line in fh:

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

            if line.startswith(">"):
                header = line.strip()[1:]
                contig_size[header] = 0

            else:
                line_len = len(line.strip())
                assembly_size += line_len
                contig_size[header] += line_len

    return assembly_size, contig_size

@MainWrapper
def main(sample_id, assembly_file, coverage_file, coverage_bp_file, bam_file,
         opts, gsize):
    """Main executor of the process_assembly_mapping template.

    Parameters
    ----------
    sample_id : str
        Sample Identification string.
    assembly_file : str
        Path to assembly file in Fasta format.
    coverage_file : str
        Path to TSV file with coverage information for each contig.
    coverage_bp_file : str
        Path to TSV file with coverage information for each base.
    bam_file : str
        Path to BAM file.
    opts : list
        List of options for processing assembly mapping.
    gsize : int
        Expected genome size

    """

    min_assembly_coverage, max_contigs = opts

    logger.info("Starting assembly mapping processing")

    # Get coverage info, total size and total coverage from the assembly
    logger.info("Parsing coverage table")
    coverage_info, a_cov = parse_coverage_table(coverage_file)
    a_size, contig_size = get_assembly_size(assembly_file)
    logger.info("Assembly processed with a total size of '{}' and coverage"
                " of '{}'".format(a_size, a_cov))
    # Get number of assembled bp after filters
    logger.info("Parsing coverage per bp table")
    coverage_bp_data = get_coverage_from_file(coverage_bp_file)

    # Assess the minimum assembly coverage
    min_coverage = evaluate_min_coverage(min_assembly_coverage, a_cov, a_size)

    # Check if filtering the assembly using the provided min_coverage will
    # reduce the final bp number to less than 80% of the estimated genome
    # size.
    # If the check below passes with True, then the filtered assembly
    # is above the 80% genome size threshold.
    filtered_assembly = "{}_filt.fasta".format(
        os.path.splitext(assembly_file)[0])
    filtered_bam = "filtered.bam"
    logger.info("Checking filtered assembly")
    if check_filtered_assembly(coverage_info, coverage_bp_data, min_coverage,
                               gsize, contig_size, int(max_contigs),
                               sample_id):
        # Filter assembly contigs based on the minimum coverage.
        logger.info("Filtered assembly passed minimum size threshold")
        logger.info("Writting filtered assembly")
        filter_assembly(assembly_file, min_coverage, coverage_info,
                        filtered_assembly)
        logger.info("Filtering BAM file according to saved contigs")
        filter_bam(coverage_info, bam_file, min_coverage, filtered_bam)
    # Could not filter the assembly as it would drop below acceptable
    # length levels. Copy the original assembly to the output assembly file
    # for compliance with the output channel
    else:
        shutil.copy(assembly_file, filtered_assembly)
        shutil.copy(bam_file, filtered_bam)
        shutil.copy(bam_file + ".bai", filtered_bam + ".bai")

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

if __name__ == '__main__':

    main(SAMPLE_ID, ASSEMBLY_FILE, COVERAGE_FILE, COVERAGE_BP_FILE,
         BAM_FILE, OPTS, GSIZE)


#!/usr/bin/env python3

"""
Purpose
-------

This module is intended to process the output of assemblies from a single
sample from programs such as Spades or Skesa.
The main input is an assembly file produced by an assembler, which will then be
filtered according to user-specified parameters.

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``: Fasta file with the assembly.
    - e.g.: ``'contigs.fasta'``
- ``opts``: List of options for processing spades assembly.
    1. Minimum contig length.
        - e.g.: ``'150'``
    2. Minimum k-mer coverage.
        - e.g.: ``'2'``
    3. Maximum number of contigs per 1.5Mb.
        - e.g.: ``'100'``
- ``assembler``: The name of the assembler
    - e.g.: ``spades``

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

(Values within ``${}`` are substituted by the corresponding variable.)

- ``'${sample_id}.assembly.fasta'`` : Fasta file with the filtered assembly.
    - e.g.: ``'Sample1.assembly.fasta'``
- ``${sample_id}.report.fasta`` : CSV file with the results of the filters for\
    each contig.
    - e.g.: ``'Sample1.report.csv'``

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

"""

__version__ = "1.0.1"
__build__ = "11042018"
__template__ = "process_assembly-nf"

import os
import json
import operator

from flowcraft_utils.flowcraft_base import get_logger, MainWrapper

logger = get_logger(__file__)

if __file__.endswith(".command.sh"):
    SAMPLE_ID = '$sample_id'
    ASSEMBLY_FILE = '$assembly'
    GSIZE = float('$gsize')
    OPTS = [x.strip() for x in '$opts'.strip("[]").split(",")]
    ASSEMBLER = '$assembler'
    logger.debug("Running {} with parameters:".format(
        os.path.basename(__file__)))
    logger.debug("SAMPLE_ID: {}".format(SAMPLE_ID))
    logger.debug("GSIZE: {}".format(GSIZE))
    logger.debug("OPTS: {}".format(OPTS))
    logger.debug("ASSEMBLER: {}".format(ASSEMBLER))

class Assembly:
    """Class that parses and filters a Fasta assembly file

    This class parses an assembly fasta file, collects a number
    of summary statistics and metadata from the contigs, filters
    contigs based on user-defined metrics and writes filtered assemblies
    and reports.

    Parameters
    ----------
    assembly_file : str
        Path to assembly file.
    min_contig_len : int
        Minimum contig length when applying the initial assembly filter.
    min_kmer_cov : int
        Minimum k-mer coverage when applying the initial assembly.
        filter.
    sample_id : str
        Name of the sample for the current assembly.
    """

    def __init__(self, assembly_file, min_contig_len, min_kmer_cov,
                 sample_id):

        self.contigs = {}
        """
        dict: Dictionary storing data for each contig.
        """

        self.filtered_ids = []
        """
        list: List of filtered contig_ids.
        """

        self.min_gc = 0.05
        """
        float: Sets the minimum GC content on a contig.
        """

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

        self.report = {}
        """
        dict: Will contain the filtering results for each contig.
        """

        self.filters = [
            ["length", ">=", min_contig_len],
            ["kmer_cov", ">=", min_kmer_cov]
        ]
        """
        list: Setting initial filters to check when parsing the assembly file.
        This can be later changed using the 'filter_contigs' method.
        """

        # Parse assembly and populate self.contigs
        self._parse_assembly(assembly_file)

        # Perform first contig filtering using min_contig_len, min_kmer_cov,
        # and gc content
        self.filter_contigs(*self.filters)

    @staticmethod
    def _parse_coverage(header_str):
        """Attempts to retrieve the coverage value from the header string.

        It splits the header by "_" and then screens the list backwards in
        search of the first float value. This will be interpreted as the
        coverage value. If it cannot find a float value, it returns None.
        This search methodology is based on the strings of assemblers
        like spades and skesa that put the mean kmer coverage for each
        contig in its corresponding fasta header.

        Parameters
        ----------
        header_str : str
            String

        Returns
        -------
        float or None
            The coverage value for the contig. None if it cannot find the
            value in the provide string.
        """

        cov = None
        for i in header_str.split("_")[::-1]:
            try:
                cov = float(i)
                break
            except ValueError:
                continue

        return cov

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

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

        The insertion of data on the self.contigs is done by the
        :py:meth:`Assembly._populate_contigs` method, which also calculates
        GC content and proportions.

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

        """

        # Temporary storage of sequence data
        seq_temp = []
        # Id counter for contig that will serve as key in self.contigs
        contig_id = 0
        # Initialize kmer coverage and header
        cov, header = None, None

        with open(assembly_file) as fh:

            logger.debug("Starting iteration of assembly file: {}".format(
                assembly_file))
            for line in fh:
                # Skip empty lines
                if not line.strip():
                    continue
                else:
                    # Remove whitespace surrounding line for further processing
                    line = line.strip()

                if line.startswith(">"):
                    # If a sequence has already been populated, save the
                    # previous contig information
                    if seq_temp:
                        # Use join() to convert string list into the full
                        # contig string. This is generally much more efficient
                        # than successively concatenating strings.
                        seq = "".join(seq_temp)

                        logger.debug("Populating contig with contig_id '{}', "
                                     "header '{}' and cov '{}'".format(
                                        contig_id, header, cov))
                        self._populate_contigs(contig_id, header, cov, seq)

                        # Reset temporary sequence storage
                        seq_temp = []
                        contig_id += 1

                    header = line[1:]
                    cov = self._parse_coverage(line)

                else:
                    seq_temp.append(line)

            # Populate last contig entry
            logger.debug("Populating contig with contig_id '{}', "
                         "header '{}' and cov '{}'".format(
                            contig_id, header, cov))
            seq = "".join(seq_temp)
            self._populate_contigs(contig_id, header, cov, seq)

    def _populate_contigs(self, contig_id, header, cov, sequence):
        """ Inserts data from a single contig into\
         :py:attr:`~Assembly.contigs`.

        By providing a contig id, the original header, the coverage that
        is parsed from the header and the sequence, this method will
        populate the :py:attr:`~Assembly.contigs` attribute.

        Parameters
        ----------
        contig_id : int
            Arbitrary unique contig identifier.
        header : str
            Original header of the current contig.
        cov : float
            The contig coverage, parsed from the fasta header
        sequence : str
            The complete sequence of the contig.

        """

        # Get AT/GC/N counts and proportions.
        # Note that self._get_gc_content returns a dictionary with the
        # information on the GC/AT/N counts and proportions. This makes it
        # much easier to add to the contigs attribute using the ** notation.
        gc_kwargs = self._get_gc_content(sequence, len(sequence))
        logger.debug("Populate GC content with: {}".format(gc_kwargs))

        self.contigs[contig_id] = {
            "header": header,
            "sequence": sequence,
            "length": len(sequence),
            "kmer_cov": cov,
            **gc_kwargs
        }

    @staticmethod
    def _get_gc_content(sequence, length):
        """Get GC content and proportions.

        Parameters
        ----------
        sequence : str
            The complete sequence of the contig.
        length : int
            The length of the sequence contig.

        Returns
        -------
        x : dict
            Dictionary with the at/gc/n counts and proportions

        """

        # Get AT/GC/N counts
        at = sum(map(sequence.count, ["A", "T"]))
        gc = sum(map(sequence.count, ["G", "C"]))
        n = length - (at + gc)

        # Get AT/GC/N proportions
        at_prop = at / length
        gc_prop = gc / length
        n_prop = n / length

        return {"at": at, "gc": gc, "n": n,
                "at_prop": at_prop, "gc_prop": gc_prop, "n_prop": n_prop}

    @staticmethod
    def _test_truth(x, op, y):
        """ Test the truth of a comparisong 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 rigth

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

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

        return ops[op](x, y)

    def filter_contigs(self, *comparisons):
        """Filters the contigs of the assembly according to user provided\
        comparisons.

        The comparisons must be a list of three elements with the
        :py:attr:`~Assembly.contigs` key, operator and test value. For
        example, to filter contigs with a minimum length of 250, a comparison
        would be::

            self.filter_contigs(["length", ">=", 250])

        The filtered contig ids will be stored in the
        :py:attr:`~Assembly.filtered_ids` list.

        The result of the test for all contigs will be stored in the
        :py:attr:`~Assembly.report` dictionary.

        Parameters
        ----------
        comparisons : list
            List with contig key, operator and value to test.

        """

        # Reset list of filtered ids
        self.filtered_ids = []
        self.report = {}

        gc_filters = [
            ["gc_prop", ">=", self.min_gc],
            ["gc_prop", "<=", 1 - self.min_gc]
        ]

        self.filters = list(comparisons) + gc_filters

        logger.debug("Filtering contigs using filters: {}".format(
            self.filters))

        for contig_id, contig in self.contigs.items():
            for key, op, value in list(comparisons) + gc_filters:
                if not self._test_truth(contig[key], op, value):
                    self.filtered_ids.append(contig_id)
                    self.report[contig_id] = "{}/{}/{}".format(key,
                                                               contig[key],
                                                               value)
                    break
                else:
                    self.report[contig_id] = "pass"

    def get_assembly_length(self):
        """Returns the length of the assembly, without the filtered contigs.

        Returns
        -------
        x : int
            Total length of the assembly.

        """

        return sum(
            [vals["length"] for contig_id, vals in self.contigs.items()
             if contig_id not in self.filtered_ids])

    def write_assembly(self, output_file, filtered=True):
        """Writes the assembly to a new file.

        The ``filtered`` option controls whether the new assembly will be
        filtered or not.

        Parameters
        ----------
        output_file : str
            Name of the output assembly file.
        filtered : bool
            If ``True``, does not include filtered ids.
        """

        logger.debug("Writing the filtered assembly into: {}".format(
            output_file))
        with open(output_file, "w") as fh:

            for contig_id, contig in self.contigs.items():
                if contig_id not in self.filtered_ids and filtered:
                    fh.write(">{}_{}\\n{}\\n".format(self.sample,
                                                     contig["header"],
                                                     contig["sequence"]))

    def write_report(self, output_file):
        """Writes a report with the test results for the current assembly

        Parameters
        ----------
        output_file : str
            Name of the output assembly file.

        """

        logger.debug("Writing the assembly report into: {}".format(
            output_file))
        with open(output_file, "w") as fh:

            for contig_id, vals in self.report.items():
                fh.write("{}, {}\\n".format(contig_id, vals))

@MainWrapper
def main(sample_id, assembly_file, gsize, opts, assembler):
    """Main executor of the process_spades template.

    Parameters
    ----------
    sample_id : str
        Sample Identification string.
    assembly_file : str
        Path to the assembly file generated by Spades.
    gsize : int
        Estimate of genome size.
    opts : list
        List of options for processing spades assembly.
    assembler : str
        Name of the assembler, for logging purposes

    """

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

    min_contig_len, min_kmer_cov, max_contigs = [int(x) for x in opts]
    logger.debug("Setting minimum conting length to: {}".format(
        min_contig_len))
    logger.debug("Setting minimum kmer coverage: {}".format(min_kmer_cov))

    # Parse the spades assembly file and perform the first filtering.
    logger.info("Starting assembly parsing")
    assembly_obj = Assembly(assembly_file, min_contig_len, min_kmer_cov,
                               sample_id)

    with open(".warnings", "w") as warn_fh:
        t_80 = gsize * 1000000 * 0.8
        t_150 = gsize * 1000000 * 1.5
        # Check if assembly size of the first assembly is lower than 80% of the
        # estimated genome size. If True, redo the filtering without the
        # k-mer coverage filter
        assembly_len = assembly_obj.get_assembly_length()
        logger.debug("Checking assembly length: {}".format(assembly_len))

        if assembly_len < t_80:

            logger.warning("Assembly size ({}) smaller than the minimum "
                           "threshold of 80% of expected genome size. "
                           "Applying contig filters without the k-mer "
                           "coverage filter".format(assembly_len))
            assembly_obj.filter_contigs(*[
                ["length", ">=", min_contig_len]
            ])

            assembly_len = assembly_obj.get_assembly_length()
            logger.debug("Checking updated assembly length: "
                         "{}".format(assembly_len))
            if assembly_len < t_80:

                warn_msg = "Assembly size smaller than the minimum" \
                           " threshold of 80% of expected genome size: {}".format(
                                assembly_len)
                logger.warning(warn_msg)
                warn_fh.write(warn_msg)
                fails = warn_msg

        if assembly_len > t_150:

            warn_msg = "Assembly size ({}) larger than the maximum" \
                       " threshold of 150% of expected genome size.".format(
                            assembly_len)
            logger.warning(warn_msg)
            warn_fh.write(warn_msg)
            fails = warn_msg

        logger.debug("Checking number of contigs: {}".format(
            len(assembly_obj.contigs)))
        contig_threshold = (max_contigs * gsize) / 1.5
        if len(assembly_obj.contigs) > contig_threshold:

            warn_msg = "The number of contigs ({}) exceeds the threshold of " \
                       "{} contigs per 1.5Mb ({})".format(
                            len(assembly_obj.contigs),
                            max_contigs,
                            round(contig_threshold, 1))

            logger.warning(warn_msg)
            warn_fh.write(warn_msg)
            warnings.append(warn_msg)

    # Write filtered assembly
    logger.debug("Renaming old assembly file to: {}".format(
        "{}.old".format(assembly_file)))
    assembly_obj.write_assembly("{}_proc.fasta".format(
        os.path.splitext(assembly_file)[0]))
    # Write report
    output_report = "{}.report.csv".format(sample_id)
    assembly_obj.write_report(output_report)
    # Write json report
    with open(".report.json", "w") as json_report:
        json_dic = {
            "tableRow": [{
                "sample": sample_id,
                "data": [
                    {"header": "Contigs ({})".format(assembler),
                     "value": len(assembly_obj.contigs),
                     "table": "assembly",
                     "columnBar": True},
                    {"header": "Assembled BP ({})".format(assembler),
                     "value": assembly_len,
                     "table": "assembly",
                     "columnBar": True}
                ]
            }],
        }

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

        if fails:
            json_dic["fail"] = [{
                "sample": sample_id,
                "table": "assembly",
                "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, ASSEMBLY_FILE, GSIZE, OPTS, ASSEMBLER)


import re
import pkgutil

try:
    from generator import components
except ImportError:
    from flowcraft.generator import components

def convert_camel_case(name):
    """Convers a CamelCase string into a snake_case one

    Parameters
    ----------
    name : str
        An arbitrary string that may be CamelCase

    Returns
    -------
    str
        The input string converted into snake_case

    """
    s1 = re.sub('(.)([A-Z][a-z]+)', r'\1_\2', name)
    return re.sub('([a-z0-9])([A-Z])', r'\1_\2', s1).lower()

def collect_process_map():
    """Collects Process classes and return dict mapping templates to classes

    This function crawls through the components module and retrieves all
    classes that inherit from the Process class. Then, it converts the name
    of the classes (which should be CamelCase) to snake_case, which is used
    as the template name.

    Returns
    -------
    dict
        Dictionary mapping the template name (snake_case) to the corresponding
        process class.
    """

    process_map = {}

    prefix = "{}.".format(components.__name__)
    for importer, modname, _ in pkgutil.iter_modules(components.__path__,
                                                     prefix):

        _module = importer.find_module(modname).load_module(modname)

        _component_classes = [
            cls for cls in _module.__dict__.values() if
            isinstance(cls, type) and cls.__name__ != "Process"
        ]

        for cls in _component_classes:
            process_map[convert_camel_case(cls.__name__)] = cls

    return process_map


#!/usr/bin/env python3

import os
import json
import dendropy

from flowcraft_utils.flowcraft_base import get_logger, MainWrapper

"""
Purpose
-------

This module is intended to process the newick generated by
 a proces to generate a report. The newick tree will be
 rooted (midpoint).

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

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

- ``newick``: phylogenetic tree in newick format.

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

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

"""

__version__ = "1.0.1"
__build__ = "20.09.2018"
__template__ = "raxml-nf"

logger = get_logger(__file__)

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

@MainWrapper
def main(newick):
    """Main executor of the process_newick template.

    Parameters
    ----------
    newick : str
        path to the newick file.

    """

    logger.info("Starting newick file processing")

    print(newick)

    tree = dendropy.Tree.get(file=open(newick, 'r'), schema="newick")

    tree.reroot_at_midpoint()

    to_write=tree.as_string("newick").strip().replace("[&R] ", '').replace(' ', '_').replace("'", "")

    with open(".report.json", "w") as json_report:
        json_dic = {
            "treeData": [{
                "trees": [
                    to_write
                ]
            }],
        }

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

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

if __name__ == '__main__':

    main(NEWICK)


#!/usr/bin/python
#
# Copyright (C) Christian Thurau, 2010.
# Licensed under the GNU General Public License (GPL).
# http://www.gnu.org/licenses/gpl.txt
"""
PyMF Compact Matrix Decomposition [1]

    CMD(CUR):  Class for Compact Matrix Decomposition

[1] Sun, J., Xie, Y., Zhang, H. and Faloutsos, C. (2007), Less is More: Compact Matrix Decomposition for Large
Sparse Graphs, in Proc. SIAM Int. Conf. on Data Mining.
"""

import numpy as np
from .cur import CUR

__all__ = ["CMD"]

class CMD(CUR):
    """
    CMD(data, rrank=0, crank=0)

    Compact Matrix Decomposition. Factorize a data matrix into three matrices s.t.
    F = | data - USV| is minimal. CMD randomly selects rows and columns from
    data for building U and V, respectively.

    Parameters
    ----------
    data : array_like [data_dimension x num_samples]
        the input data
    rrank: int, optional
        Number of rows to sample from data. Double entries are eliminiated s.t.
        the resulting rank might be lower.
        4 (default)
    crank: int, optional
        Number of columns to sample from data. Double entries are eliminiated s.t.
        the resulting rank might be lower.
        4 (default)

    Attributes
    ----------
        U,S,V : submatrices s.t. data = USV

    Example
    -------
    >>> import numpy as np
    >>> from cmd import CMD
    >>> data = np.array([[1.0, 0.0, 2.0], [0.0, 1.0, 1.0]])
    >>> cmd_mdl = CMD(data, show_progress=False, rrank=1, crank=2)
    >>> cmd_mdl.factorize()
    """

    def _cmdinit(self):
        nrids = np.unique(self._rid)
        ncids = np.unique(self._cid)

        self._rcnt = np.zeros(len(nrids))
        self._ccnt = np.zeros(len(ncids))

        for i,idx in enumerate(nrids):
            self._rcnt[i] = len(np.where(self._rid == idx)[0])

        for i,idx in enumerate(ncids):
            self._ccnt[i] = len(np.where(self._cid == idx)[0])

        self._rid = np.int32(list(nrids))
        self._cid = np.int32(list(ncids))

    def factorize(self):
        """ Factorize s.t. CUR = data

            Updated Values
            --------------
            .C : updated values for C.
            .U : updated values for U.
            .R : updated values for R.
        """

        [prow, pcol] = self.sample_probability()

        self._rid = self.sample(self._rrank, prow)
        self._cid = self.sample(self._crank, pcol)

        self._cmdinit()

        self.computeUCR()

if __name__ == "__main__":
    import doctest
    doctest.testmod()


#!/usr/bin/python
#
# Copyright (C) Christian Thurau, 2010.
# Licensed under the GNU General Public License (GPL).
# http://www.gnu.org/licenses/gpl.txt
"""
PyMF CUR Decomposition [1]

    CUR(SVD) : Class for CUR Decomposition

[1] Drineas, P., Kannan, R. and Mahoney, M. (2006), 'Fast Monte Carlo Algorithms III: Computing
a Compressed Approixmate Matrix Decomposition', SIAM J. Computing 36(1), 184-206.
"""

import numpy as np
import scipy.sparse

from .svd import pinv, SVD

__all__ = ["CUR"]

class CUR(SVD):
    """
    CUR(data,  data, k=-1, rrank=0, crank=0)

    CUR Decomposition. Factorize a data matrix into three matrices s.t.
    F = | data - USV| is minimal. CUR randomly selects rows and columns from
    data for building U and V, respectively.

    Parameters
    ----------
    data : array_like [data_dimension x num_samples]
        the input data
    rrank: int, optional
        Number of rows to sample from data.
        4 (default)
    crank: int, optional
        Number of columns to sample from data.
        4 (default)
    show_progress: bool, optional
        Print some extra information
        False (default)

    Attributes
    ----------
        U,S,V : submatrices s.t. data = USV

    Example
    -------
    >>> import numpy as np
    >>> from cur import CUR
    >>> data = np.array([[1.0, 0.0, 2.0], [0.0, 1.0, 1.0]])
    >>> cur_mdl = CUR(data, show_progress=False, rrank=1, crank=2)
    >>> cur_mdl.factorize()
    """

    def __init__(self, data, k=-1, rrank=0, crank=0):
        SVD.__init__(self, data,k=k,rrank=rrank, crank=rrank)

        # select all data samples for computing the error:
        # note that this might take very long, adjust self._rset and self._cset
        # for faster computations.
        self._rset = range(self._rows)
        self._cset = range(self._cols)

    def sample(self, s, probs):
        prob_rows = np.cumsum(probs.flatten())
        temp_ind = np.zeros(s, np.int32)

        for i in range(s):
            v = np.random.rand()

            try:
                tempI = np.where(prob_rows >= v)[0]
                temp_ind[i] = tempI[0]
            except:
                temp_ind[i] = len(prob_rows)

        return np.sort(temp_ind)

    def sample_probability(self):

        if scipy.sparse.issparse(self.data):
            dsquare = self.data.multiply(self.data)
        else:
            dsquare = self.data[:,:]**2

        prow = np.array(dsquare.sum(axis=1), np.float64)
        pcol = np.array(dsquare.sum(axis=0), np.float64)

        prow /= prow.sum()
        pcol /= pcol.sum()

        return (prow.reshape(-1,1), pcol.reshape(-1,1))

    def computeUCR(self):
        # the next  lines do NOT work with h5py if CUR is used -> double indices in self.cid or self.rid
        # can occur and are not supported by h5py. When using h5py data, always use CMD which ignores
        # reoccuring row/column selections.

        if scipy.sparse.issparse(self.data):
            self._C = self.data[:, self._cid] * scipy.sparse.csc_matrix(np.diag(self._ccnt**(1/2)))
            self._R = scipy.sparse.csc_matrix(np.diag(self._rcnt**(1/2))) * self.data[self._rid,:]

            self._U = pinv(self._C, self._k) * self.data[:,:] * pinv(self._R, self._k)

        else:
            self._C = np.dot(self.data[:, self._cid].reshape((self._rows, len(self._cid))), np.diag(self._ccnt**(1/2)))
            self._R = np.dot(np.diag(self._rcnt**(1/2)), self.data[self._rid,:].reshape((len(self._rid), self._cols)))

            self._U = np.dot(np.dot(pinv(self._C, self._k), self.data[:,:]),
                             pinv(self._R, self._k))

        # set some standard (with respect to SVD) variable names
        self.U = self._C
        self.S = self._U
        self.V = self._R

    def factorize(self):
        """ Factorize s.t. CUR = data

            Updated Values
            --------------
            .C : updated values for C.
            .U : updated values for U.
            .R : updated values for R.
        """
        [prow, pcol] = self.sample_probability()
        self._rid = self.sample(self._rrank, prow)
        self._cid = self.sample(self._crank, pcol)

        self._rcnt = np.ones(len(self._rid))
        self._ccnt = np.ones(len(self._cid))

        self.computeUCR()

if __name__ == "__main__":
    import doctest
    doctest.testmod()


#!/usr/bin/python
#
# Copyright (C) Christian Thurau, 2010.
# Licensed under the GNU General Public License (GPL).
# http://www.gnu.org/licenses/gpl.txt
"""
PyMF Convex Hull Non-negative Matrix Factorization [1]

    CHNMF(NMF) : Class for Convex-hull NMF
    quickhull : Function for finding the convex hull in 2D

[1] C. Thurau, K. Kersting, and C. Bauckhage. Convex Non-Negative Matrix
Factorization in the Wild. ICDM 2009.
"""

import numpy as np

from itertools import combinations
from .dist import vq
from .pca import PCA
from .aa import AA

__all__ = ["CHNMF"]

def quickhull(sample):
    """ Find data points on the convex hull of a supplied data set

    Args:
        sample: data points as column vectors n x d
                    n - number samples
                    d - data dimension (should be two)

    Returns:
        a k x d matrix containint the convex hull data points
    """

    link = lambda a, b: np.concatenate((a, b[1:]))
    edge = lambda a, b: np.concatenate(([a], [b]))

    def dome(sample, base):
        h, t = base
        dists = np.dot(sample - h, np.dot(((0, -1), (1, 0)), (t - h)))
        outer = np.repeat(sample, dists > 0, axis=0)

        if len(outer):
            pivot = sample[np.argmax(dists)]
            return link(dome(outer, edge(h, pivot)),
                dome(outer, edge(pivot, t)))
        else:
            return base

    if len(sample) > 2:
        axis = sample[:, 0]
        base = np.take(sample, [np.argmin(axis), np.argmax(axis)], axis=0)
        return link(dome(sample, base),
            dome(sample, base[::-1]))
    else:
        return sample

class CHNMF(AA):
    """
    CHNMF(data, num_bases=4)

    Convex Hull Non-negative Matrix Factorization. Factorize a data matrix into
    two matrices s.t. F = | data - W*H | is minimal. H is restricted to convexity
    (H >=0, sum(H, axis=1) = [1 .. 1]) and W resides on actual data points.
    Factorization is solved via an alternating least squares optimization using
    the quadratic programming solver from cvxopt. The results are usually
    equivalent to Archetypal Analysis (pymf.AA) but CHNMF also works for very
    large datasets.

    Parameters
    ----------
    data : array_like, shape (_data_dimension, _num_samples)
        the input data
    num_bases: int, optional
        Number of bases to compute (column rank of W and row rank of H).
        4 (default)
    base_sel: int,
        Number of pairwise basis vector projections. Set to a value< rank(data).
        Computation time scale exponentially with this value, usually rather low
        values are sufficient (3-10).

    Attributes
    ----------
        W : "data_dimension x num_bases" matrix of basis vectors
        H : "num bases x num_samples" matrix of coefficients
        ferr : frobenius norm (after calling .factorize())

    Example
    -------
    Applying CHNMF to some rather stupid data set:

    >>> import numpy as np
    >>> from chnmf import CHNMF
    >>> data = np.array([[1.0, 0.0, 2.0], [0.0, 1.0, 1.0]])

    Use 2 basis vectors -> W shape(data_dimension, 2).

    >>> chnmf_mdl = CHNMF(data, num_bases=2)

    And start computing the factorization.

    >>> chnmf_mdl.factorize()

    The basis vectors are now stored in chnmf_mdl.W, the coefficients in
    chnmf_mdl.H. To compute coefficients for an existing set of basis vectors
    simply copy W to chnmf_mdl.W, and set compute_w to False:

    >>> data = np.array([[1.5, 2.0], [1.2, 1.8]])
    >>> W = np.array([[1.0, 0.0], [0.0, 1.0]])
    >>> chnmf_mdl = CHNMF(data, num_bases=2)
    >>> chnmf_mdl.W = W
    >>> chnmf_mdl.factorize(compute_w=False)

    The result is a set of coefficients chnmf_mdl.H, s.t. data = W * chnmf_mdl.H.
    """

    def __init__(self, data, num_bases=4, base_sel=3):

        # call inherited method
        AA.__init__(self, data, num_bases=num_bases)

        # base sel should never be larger than the actual data dimension
        self._base_sel = base_sel
        if base_sel > self.data.shape[0]:
            self._base_sel = self.data.shape[0]

    def init_h(self):
        self.H = np.zeros((self._num_bases, self._num_samples))

    def init_w(self):
        self.W = np.zeros((self._data_dimension, self._num_bases))

    def _map_w_to_data(self):
        """ Return data points that are most similar to basis vectors W
        """

        # assign W to the next best data sample
        self._Wmapped_index = vq(self.data, self.W)
        self.Wmapped = np.zeros(self.W.shape)

        # do not directly assign, i.e. Wdist = self.data[:,sel]
        # as self might be unsorted (in non ascending order)
        # -> sorting sel would screw the matching to W if
        # self.data is stored as a hdf5 table (see h5py)
        for i, s in enumerate(self._Wmapped_index):
            self.Wmapped[:,i] = self.data[:,s]

    def update_w(self):
        """ compute new W """
        def select_hull_points(data, n=3):
            """ select data points for pairwise projections of the first n
            dimensions """

            # iterate over all projections and select data points
            idx = np.array([])

            # iterate over some pairwise combinations of dimensions
            for i in combinations(range(n), 2):
                # sample convex hull points in 2D projection
                convex_hull_d = quickhull(data[i, :].T)

                # get indices for convex hull data points
                idx = np.append(idx, vq(data[i, :], convex_hull_d.T))
                idx = np.unique(idx)

            return np.int32(idx)

        # determine convex hull data points using either PCA or random
        # projections
        method = 'randomprojection'
        if method == 'pca':
            pcamodel = PCA(self.data)
            pcamodel.factorize(show_progress=False)
            proj = pcamodel.H
        else:
            R = np.random.randn(self._base_sel, self._data_dimension)
            proj = np.dot(R, self.data)

        self._hull_idx = select_hull_points(proj, n=self._base_sel)
        aa_mdl = AA(self.data[:, self._hull_idx], num_bases=self._num_bases)

        # determine W
        aa_mdl.factorize(niter=50, compute_h=True, compute_w=True,
                         compute_err=True, show_progress=False)

        self.W = aa_mdl.W
        self._map_w_to_data()

    def factorize(self, show_progress=False, compute_w=True, compute_h=True,
                  compute_err=True, niter=1):
        """ Factorize s.t. WH = data

            Parameters
            ----------
            show_progress : bool
                    print some extra information to stdout.
            compute_h : bool
                    iteratively update values for H.
            compute_w : bool
                    iteratively update values for W.
            compute_err : bool
                    compute Frobenius norm |data-WH| after each update and store
                    it to .ferr[k].

            Updated Values
            --------------
            .W : updated values for W.
            .H : updated values for H.
            .ferr : Frobenius norm |data-WH|.
        """

        AA.factorize(self, niter=1, show_progress=show_progress,
                  compute_w=compute_w, compute_h=compute_h,
                  compute_err=compute_err)

if __name__ == "__main__":
    import doctest
    doctest.testmod()


#!/usr/bin/python
#
# Copyright (C) Christian Thurau, 2010.
# Licensed under the GNU General Public License (GPL).
# http://www.gnu.org/licenses/gpl.txt
"""
PyMF Geometric-Map

    GMAP: Class for Geometric-Map
"""

import scipy.sparse
import numpy as np

from .dist import *
from .aa import AA
from .kmeans import Kmeans

__all__ = ["GMAP"]

class GMAP(AA):
    """
    GMAP(data, num_bases=4, dist_measure='l2')

    Geometric-Map. Factorize a data matrix into two matrices s.t.
    F = | data - W*H | is minimal. G-MAP can emulate/approximate several
    standard methods including PCA, NMF, and AA.

    Parameters
    ----------
    data : array_like, shape (_data_dimension, _num_samples)
        the input data
    num_bases: int, optional
        Number of bases to compute (column rank of W and row rank of H).
        4 (default)
    method : one of 'pca' ,'nmf', 'aa', default is 'pca' which emulates
        Principal Component Analysis using the geometric map method ('nmf'
        emulates Non-negative Matrix Factorization, 'aa' emulates Archetypal
        Analysis).
    robust_map : bool, optional
        use robust_map or the standard max-val selection
        [see "On FastMap and the Convex Hull of Multivariate Data: Toward
        Fast and Robust Dimension Reduction", Ostrouchov and Samatova, PAMI
        2005]
    Attributes
    ----------
    W : "data_dimension x num_bases" matrix of basis vectors
    H : "num bases x num_samples" matrix of coefficients
    ferr : frobenius norm (after calling .factorize())

    Example
    -------
    Applying GMAP to some rather stupid data set:

    >>> import numpy as np
    >>> data = np.array([[1.0, 0.0, 2.0], [0.0, 1.0, 1.0]])
    >>> gmap_mdl = GMAP(data, num_bases=2)
    >>> gmap_mdl.factorize()

    The basis vectors are now stored in gmap_mdl.W, the coefficients in gmap_mdl.H.
    To compute coefficients for an existing set of basis vectors simply copy W
    to gmap_mdl.W, and set compute_w to False:

    >>> data = np.array([[1.5, 1.3], [1.2, 0.3]])
    >>> W = np.array([[1.0, 0.0], [0.0, 1.0]])
    >>> gmap_mdl = GMAP(data, num_bases=2)
    >>> gmap_mdl.W = W
    >>> gmap_mdl.factorize(compute_w=False)

    The result is a set of coefficients gmap_mdl.H, s.t. data = W * gmap_mdl.H.
    """

    # always overwrite the default number of iterations
    # -> any value other does not make sense.
    _NITER = 1

    def __init__(self, data, num_bases=4, method='pca', robust_map=True):

        AA.__init__(self, data, num_bases=num_bases)
        self.sub = []
        self._robust_map = robust_map
        self._method = method

    def init_h(self):
        self.H = np.zeros((self._num_bases, self._num_samples))

    def init_w(self):
        self.W = np.zeros((self._data_dimension, self._num_bases))

    def update_w(self):
        """ compute new W """

        def select_next(iterval):
            """ select the next best data sample using robust map
            or simply the max iterval ... """

            if self._robust_map:
                k = np.argsort(iterval)[::-1]
                d_sub = self.data[:,k[:self._robust_nselect]]
                self.sub.extend(k[:self._robust_nselect])

                # cluster d_sub
                kmeans_mdl = Kmeans(d_sub, num_bases=self._robust_cluster)
                kmeans_mdl.factorize(niter=10)

                # get largest cluster
                h = np.histogram(kmeans_mdl.assigned, range(self._robust_cluster+1))[0]
                largest_cluster = np.argmax(h)
                sel = pdist(kmeans_mdl.W[:, largest_cluster:largest_cluster+1], d_sub)
                sel = k[np.argmin(sel)]
            else:
                sel = np.argmax(iterval)

            return sel

        EPS = 10**-8

        if scipy.sparse.issparse(self.data):
            norm_data = np.sqrt(self.data.multiply(self.data).sum(axis=0))
            norm_data = np.array(norm_data).reshape((-1))
        else:
            norm_data = np.sqrt(np.sum(self.data**2, axis=0))

        self.select = []

        if self._method == 'pca' or self._method == 'aa':
            iterval = norm_data.copy()

        if self._method == 'nmf':
            iterval = np.sum(self.data, axis=0)/(np.sqrt(self.data.shape[0])*norm_data)
            iterval = 1.0 - iterval

        self.select.append(select_next(iterval))

        for l in range(1, self._num_bases):

            if scipy.sparse.issparse(self.data):
                c = self.data[:, self.select[-1]:self.select[-1]+1].T * self.data
                c = np.array(c.todense())
            else:
                c = np.dot(self.data[:,self.select[-1]], self.data)

            c = c/(norm_data * norm_data[self.select[-1]])

            if self._method == 'pca':
                c = 1.0 - np.abs(c)
                c = c * norm_data

            elif self._method == 'aa':
                c = (c*-1.0 + 1.0)/2.0
                c = c * norm_data

            elif self._method == 'nmf':
                c = 1.0 - np.abs(c)

            ### update the estimated volume
            iterval = c * iterval

            # detect the next best data point
            self.select.append(select_next(iterval))

            self._logger.info('cur_nodes: ' + str(self.select))

        # sort indices, otherwise h5py won't work
        self.W = self.data[:, np.sort(self.select)]

        # "unsort" it again to keep the correct order
        self.W = self.W[:, np.argsort(np.argsort(self.select))]

    def factorize(self, show_progress=False, compute_w=True, compute_h=True,
                  compute_err=True, robust_cluster=3, niter=1, robust_nselect=-1):
        """ Factorize s.t. WH = data

            Parameters
            ----------
            show_progress : bool
                    print some extra information to stdout.
                    False, default
            compute_h : bool
                    iteratively update values for H.
                    True, default
            compute_w : bool
                    iteratively update values for W.
                    default, True
            compute_err : bool
                    compute Frobenius norm |data-WH| after each update and store
                    it to .ferr[k].
            robust_cluster : int, optional
                    set the number of clusters for robust map selection.
                    3, default
            robust_nselect : int, optional
                    set the number of samples to consider for robust map
                    selection.
                    -1, default (automatically determine suitable number)

            Updated Values
            --------------
            .W : updated values for W.
            .H : updated values for H.
            .ferr : Frobenius norm |data-WH|.
        """
        self._robust_cluster = robust_cluster
        self._robust_nselect = robust_nselect

        if self._robust_nselect == -1:
            self._robust_nselect = np.round(np.log(self.data.shape[1])*2)

        AA.factorize(self, niter=1, show_progress=show_progress,
                  compute_w=compute_w, compute_h=compute_h,
                  compute_err=compute_err)

if __name__ == "__main__":
    import doctest
    doctest.testmod()


#!/usr/bin/env python
# coding: utf-8
import numpy as np

from msaf.algorithms.interface import SegmenterInterface
from msaf.base import Features
from . import main2

class Segmenter(SegmenterInterface):
    """
    This script identifies the boundaries of a given track using the Spectral
    Clustering method published here:

    Mcfee, B., & Ellis, D. P. W. (2014). Analyzing Song Structure with Spectral
    Clustering. In Proc. of the 15th International Society for Music
    Information Retrieval Conference (pp. 405–410). Taipei, Taiwan.

    Original code by Brian McFee from:
        https://github.com/bmcfee/laplacian_segmentation
    """
    def process(self):
        """Main process.
        Returns
        -------
        est_idxs : np.array(N) or list
            Estimated times for the segment boundaries in frame indeces.
            List if hierarchical segmentation.
        est_labels : np.array(N-1) or list
            Estimated labels for the segments.
            List if hierarchical segmentation.
        """
        # This algorithm only accepts one specific kind of features:
        # Combination of PCP + MFCC. Let's get them:
        pcp_obj = Features.select_features(
            "pcp", self.file_struct, self.annot_beats, self.framesync)
        mfcc_obj = Features.select_features(
            "mfcc", self.file_struct, self.annot_beats, self.framesync)

        # Get frame times and make sure they're the same in both features
        frame_times = pcp_obj.frame_times
        assert np.array_equal(frame_times, mfcc_obj.frame_times)

        # Brian wants PCP and MFCC
        # (tranpsosed, because he's that kind of person)
        # TODO: self.in_bound_idxs
        est_idxs, est_labels, F = main2.do_segmentation(
            pcp_obj.features.T, mfcc_obj.features.T, self.config,
            self.in_bound_idxs)

        return est_idxs, est_labels, F

    def processFlat(self):
        """Main process.for flat segmentation.
        Returns
        -------
        est_idxs : np.array(N)
            Estimated times for the segment boundaries in frame indeces.
        est_labels : np.array(N-1)
            Estimated labels for the segments.
        """
        self.config["hier"] = False
        est_idxs, est_labels, F = self.process()
        assert est_idxs[0] == 0 and est_idxs[-1] == F.shape[1] - 1
        return self._postprocess(est_idxs, est_labels)

    def processHierarchical(self):
        """Main process.for hierarchial segmentation.
        Returns
        -------
        est_idxs : list
            List with np.arrays for each layer of segmentation containing
            the estimated indeces for the segment boundaries.
        est_labels : list
            List with np.arrays containing the labels for each layer of the
            hierarchical segmentation.
        """
        self.config["hier"] = True
        est_idxs, est_labels, F = self.process()
        for layer in range(len(est_idxs)):
            assert est_idxs[layer][0] == 0 and \
                est_idxs[layer][-1] == F.shape[1] - 1
            est_idxs[layer], est_labels[layer] = \
                self._postprocess(est_idxs[layer], est_labels[layer])
        return est_idxs, est_labels