<commit_before>"""
Generate a hints database file from RNAseq alignments for AugustusTMR/AugustusCGP.

Expects a config file to be passed in with the paths to the files. Example:

[ANNOTATION]
annotation = /path/to/gff3

[INTRONBAM]
genome1 = /path/to/non_polyA_bam1.bam, /path/to/non_polyA_bam2.bam

[BAM]
genome1 = /path/to/fofn

The annotation field is optional, but will help AugustusCGP make better predictions.

BAM annotations can be put either under INTRONBAM or BAM. Any INTRONBAM will only have intron data loaded, and is
suitable for lower quality RNA-seq.

"""
import collections
import itertools
import logging
import os
import sqlite3
import argparse

import luigi
import luigi.contrib.sqla
import pyfasta
import pysam
from configobj import ConfigObj
from luigi.util import requires
from toil.common import Toil
from toil.job import Job

import tools.dataOps
import tools.fileOps
import tools.mathOps
import tools.misc
import tools.procOps
import tools.toilInterface
import tools.transcripts
import tools.hal
from base_tasks import HintsDbToilTask, HintsDbTask, HintsDbWrapperTask

logger = logging.getLogger(__name__)


class UserException(Exception):
    pass


class MissingFileException(UserException):
    pass


class BuildHints(HintsDbWrapperTask):
    """
    Main entry point. Parses input files, and launches the next task.
    """
    def parse_cfg(self):
        # configspec validates the input config file
        configspec = ['[ANNOTATION]', '__many__ = string', '[INTRONBAM]', '__many__ = list', '[BAM]', '__many__ = list']
        parser = ConfigObj(self.config, configspec=configspec)

        # convert the config into a new dict, parsing the BAMs
        cfg = collections.defaultdict(dict)
        target_genomes = set()
        if 'ANNOTATION' not in parser:
            cfg['ANNOTATION'] = {}
        else:
            for genome, annot in parser['ANNOTATION'].iteritems():
                annot = os.path.abspath(annot)
                if not os.path.exists(annot):
                    raise MissingFileException('Missing annotation file {}.'.format(annot))
                cfg['ANNOTATION'][genome] = annot
                target_genomes.add(genome)

        # if a given genome only has one BAM, it is a string. Fix this. Extract all paths from fofn files.
        for dtype in ['BAM', 'INTRONBAM']:
            if dtype not in parser:  # the user does not have to specify all field types
                cfg[dtype] = {}
                continue
            for genome in parser[dtype]:
                target_genomes.add(genome)
                path = parser[dtype][genome]
                if isinstance(path, str):
                    if not tools.misc.is_bam(path):
                        # this is a fofn
                        cfg[dtype][genome] = [os.path.abspath(x.rstrip()) for x in open(path)]
                    else:
                        # this is a single BAM
                        cfg[dtype][genome] = [os.path.abspath(path)]
                else:
                    cfg[dtype][genome] = [os.path.abspath(x) for x in path]
        return cfg, tuple(target_genomes)

    def validate(self, cfg, genomes, target_genomes):
        for dtype in ['BAM', 'INTRONBAM']:
            for genome in cfg[dtype]:
                for bam in cfg[dtype][genome]:
                    if not os.path.exists(bam):
                        raise MissingFileException('Missing BAM {}.'.format(bam))
                    if not os.path.exists(bam + '.bai'):
                        raise MissingFileException('Missing BAM index {}.'.format(bam + '.bai'))

        for genome, annot in cfg['ANNOTATION'].iteritems():
            if not os.path.exists(annot):
                raise MissingFileException('Missing annotation file {}.'.format(annot))

        if all(g in genomes for g in target_genomes) is False:
            bad_genomes = set(genomes) - set(target_genomes)
            err_msg = 'Genomes {} present in configuration and not present in HAL'.format(','.join(bad_genomes))
            raise UserException(err_msg)

    def requires(self):
        hal = os.path.abspath(self.hal)
        genomes = tools.hal.extract_genomes(hal)
        cfg, target_genomes = self.parse_cfg()
        self.validate(cfg, genomes, target_genomes)
        hint_paths = {}
        flat_fasta_paths = {}
        for genome in genomes:
            flat_fasta = self.clone(GenomeFlatFasta, genome=genome, cfg=cfg, hal=hal)
            yield flat_fasta
            flat_fasta_paths[genome] = flat_fasta.output().path
            if genome in target_genomes:
                annotation = cfg['ANNOTATION'].get(genome, None)
                hints = self.clone(GenerateHints, genome=genome, flat_fasta=flat_fasta.output().path,
                                   annotation=annotation, cfg=cfg)
                yield hints
                hint_paths[genome] = hints.output().path
        yield self.clone(BuildDb, cfg=cfg, hint_paths=hint_paths, flat_fasta_paths=flat_fasta_paths, genomes=genomes,
                         target_genomes=target_genomes)


class GenomeFlatFasta(HintsDbTask):
    """
    Flattens a genome fasta using pyfasta, copying it to the work directory. Requires the pyfasta package.
    """
    genome = luigi.Parameter()
    cfg = luigi.Parameter()
    hal = luigi.Parameter()

    def output(self):
        path = os.path.abspath(os.path.join(self.work_dir, self.genome + '.fa'))
        tools.fileOps.ensure_file_dir(path)
        return luigi.LocalTarget(path)

    def run(self):
        logger.info('Extracting fasta for {} from hal.'.format(self.genome))
        with self.output().open('w') as outf:
            cmd = ['hal2fasta', self.hal, self.genome]
            tools.procOps.run_proc(cmd, stdout=outf)
        logger.info('Flattening fasta for {}.'.format(self.genome))
        cmd = ['pyfasta', 'flatten', self.output().path]
        tools.procOps.run_proc(cmd)


@requires(GenomeFlatFasta)
class GenerateHints(HintsDbToilTask):
    """
    Generate hints for each genome as a separate Toil pipeline.
    """
    genome = luigi.Parameter()
    flat_fasta = luigi.Parameter()
    annotation = luigi.Parameter()
    cfg = luigi.Parameter()

    def output(self):
        hints = os.path.abspath(os.path.join(self.work_dir, self.genome + '.extrinsic_hints.gff'))
        return luigi.LocalTarget(hints)

    def run(self):
        logger.info('Beginning GenerateHints Toil pipeline for {}.'.format(self.genome))
        work_dir = os.path.abspath(os.path.join(self.work_dir, 'toil', self.genome))
        toil_options = self.prepare_toil_options(work_dir)
        toil_options.defaultMemory = '8G'  # TODO: don't hardcode this.
        generate_hints(self.genome, self.flat_fasta, self.cfg, self.annotation, self.output().path, toil_options)
        logger.info('Finished GenerateHints Toil pipeline for {}.'.format(self.genome))


class BuildDb(HintsDbTask):
    """
    Constructs the hints database from a series of reduced hints GFFs

    TODO: output() should be way smarter than this. Currently, it only checks if the indices have been created.
    """
    cfg = luigi.Parameter()
    hint_paths = luigi.Parameter()
    flat_fasta_paths = luigi.Parameter()
    genomes = luigi.TupleParameter()
    target_genomes = luigi.TupleParameter()

    def requires(self):
        for genome in self.target_genomes:
            flat_fasta = self.clone(GenomeFlatFasta, genome=genome, cfg=self.cfg)
            annotation = self.cfg['ANNOTATION'].get(genome, None)
            hints = self.clone(GenerateHints, genome=genome, flat_fasta=flat_fasta.output().path,
                               annotation=annotation, cfg=self.cfg)
            yield hints

    def output(self):
        tools.fileOps.ensure_file_dir(self.augustus_hints_db)
        return IndexTarget(self.augustus_hints_db)

    def run(self):
        for genome in self.genomes:
            logger.info('Loading sequence for {} into database.'.format(genome))
            base_cmd = ['load2sqlitedb', '--noIdx', '--clean', '--species={}'.format(genome),
                        '--dbaccess={}'.format(self.augustus_hints_db)]
            tools.procOps.run_proc(base_cmd + [self.flat_fasta_paths[genome]])
            if genome in self.hint_paths:
                logger.info('Loading hints for {} into database.'.format(genome))
                tools.procOps.run_proc(base_cmd + [self.hint_paths[genome]])
        logger.info('Indexing database.')
        cmd = ['load2sqlitedb', '--makeIdx', '--clean', '--dbaccess={}'.format(self.augustus_hints_db)]
        tools.procOps.run_proc(cmd)


class IndexTarget(luigi.Target):<commit_msg>More error checking for hints database construction.<commit_after>"""
Generate a hints database file from RNAseq alignments for AugustusTMR/AugustusCGP.

Expects a config file to be passed in with the paths to the files. Example:

[ANNOTATION]
annotation = /path/to/gff3

[INTRONBAM]
genome1 = /path/to/non_polyA_bam1.bam, /path/to/non_polyA_bam2.bam

[BAM]
genome1 = /path/to/fofn

The annotation field is optional, but will help AugustusCGP make better predictions.

BAM annotations can be put either under INTRONBAM or BAM. Any INTRONBAM will only have intron data loaded, and is
suitable for lower quality RNA-seq.

"""
import collections
import itertools
import logging
import os
import sqlite3
import argparse

import luigi
import luigi.contrib.sqla
import pyfasta
import pysam
from configobj import ConfigObj
from luigi.util import requires
from toil.common import Toil
from toil.job import Job

import tools.dataOps
import tools.fileOps
import tools.mathOps
import tools.misc
import tools.procOps
import tools.toilInterface
import tools.transcripts
import tools.hal
from base_tasks import HintsDbToilTask, HintsDbTask, HintsDbWrapperTask

logger = logging.getLogger(__name__)


class UserException(Exception):
    pass


class MissingFileException(UserException):
    pass


class BuildHints(HintsDbWrapperTask):
    """
    Main entry point. Parses input files, and launches the next task.
    """
    def parse_cfg(self):
        if not os.path.exists(self.config):
            raise MissingFileException('Config file {} not found.'.format(self.config))
        # configspec validates the input config file
        configspec = ['[ANNOTATION]', '__many__ = string', '[INTRONBAM]', '__many__ = list', '[BAM]', '__many__ = list']
        parser = ConfigObj(self.config, configspec=configspec)

        # convert the config into a new dict, parsing the BAMs
        cfg = collections.defaultdict(dict)
        target_genomes = set()
        if 'ANNOTATION' not in parser:
            cfg['ANNOTATION'] = {}
        else:
            for genome, annot in parser['ANNOTATION'].iteritems():
                annot = os.path.abspath(annot)
                if not os.path.exists(annot):
                    raise MissingFileException('Missing annotation file {}.'.format(annot))
                cfg['ANNOTATION'][genome] = annot
                target_genomes.add(genome)

        # if a given genome only has one BAM, it is a string. Fix this. Extract all paths from fofn files.
        for dtype in ['BAM', 'INTRONBAM']:
            if dtype not in parser:  # the user does not have to specify all field types
                cfg[dtype] = {}
                continue
            for genome in parser[dtype]:
                target_genomes.add(genome)
                path = parser[dtype][genome]
                if isinstance(path, str):
                    if not tools.misc.is_bam(path):
                        # this is a fofn
                        cfg[dtype][genome] = [os.path.abspath(x.rstrip()) for x in open(path)]
                    else:
                        # this is a single BAM
                        cfg[dtype][genome] = [os.path.abspath(path)]
                else:
                    cfg[dtype][genome] = [os.path.abspath(x) for x in path]
        return cfg, tuple(target_genomes)

    def validate(self, cfg, genomes, target_genomes):
        if len(cfg['BAM']) + len(cfg['INTRONBAM']) + len(cfg['ANNOTATION']) == 0:
            logger.warning('No RNA-seq or annotations found in config. Will load genomes only.')
        elif len(cfg['BAM']) + len(cfg['INTRONBAM']):
            logger.warning('No RNA-seq found in config. Will load genomes and annotation only.')

        for dtype in ['BAM', 'INTRONBAM']:
            for genome in cfg[dtype]:
                for bam in cfg[dtype][genome]:
                    if not os.path.exists(bam):
                        raise MissingFileException('Missing BAM {}.'.format(bam))
                    if not os.path.exists(bam + '.bai'):
                        raise MissingFileException('Missing BAM index {}.'.format(bam + '.bai'))

        for genome, annot in cfg['ANNOTATION'].iteritems():
            if not os.path.exists(annot):
                raise MissingFileException('Missing annotation file {}.'.format(annot))

        if all(g in genomes for g in target_genomes) is False:
            bad_genomes = set(genomes) - set(target_genomes)
            err_msg = 'Genomes {} present in configuration and not present in HAL'.format(','.join(bad_genomes))
            raise UserException(err_msg)

    def requires(self):
        hal = os.path.abspath(self.hal)
        genomes = tools.hal.extract_genomes(hal)
        cfg, target_genomes = self.parse_cfg()
        self.validate(cfg, genomes, target_genomes)
        hint_paths = {}
        flat_fasta_paths = {}
        for genome in genomes:
            flat_fasta = self.clone(GenomeFlatFasta, genome=genome, cfg=cfg, hal=hal)
            yield flat_fasta
            flat_fasta_paths[genome] = flat_fasta.output().path
            if genome in target_genomes:
                annotation = cfg['ANNOTATION'].get(genome, None)
                hints = self.clone(GenerateHints, genome=genome, flat_fasta=flat_fasta.output().path,
                                   annotation=annotation, cfg=cfg)
                yield hints
                hint_paths[genome] = hints.output().path
        yield self.clone(BuildDb, cfg=cfg, hint_paths=hint_paths, flat_fasta_paths=flat_fasta_paths, genomes=genomes,
                         target_genomes=target_genomes)


class GenomeFlatFasta(HintsDbTask):
    """
    Flattens a genome fasta using pyfasta, copying it to the work directory. Requires the pyfasta package.
    """
    genome = luigi.Parameter()
    cfg = luigi.Parameter()
    hal = luigi.Parameter()

    def output(self):
        path = os.path.abspath(os.path.join(self.work_dir, self.genome + '.fa'))
        tools.fileOps.ensure_file_dir(path)
        return luigi.LocalTarget(path)

    def run(self):
        logger.info('Extracting fasta for {} from hal.'.format(self.genome))
        with self.output().open('w') as outf:
            cmd = ['hal2fasta', self.hal, self.genome]
            tools.procOps.run_proc(cmd, stdout=outf)
        logger.info('Flattening fasta for {}.'.format(self.genome))
        cmd = ['pyfasta', 'flatten', self.output().path]
        tools.procOps.run_proc(cmd)


@requires(GenomeFlatFasta)
class GenerateHints(HintsDbToilTask):
    """
    Generate hints for each genome as a separate Toil pipeline.
    """
    genome = luigi.Parameter()
    flat_fasta = luigi.Parameter()
    annotation = luigi.Parameter()
    cfg = luigi.Parameter()

    def output(self):
        hints = os.path.abspath(os.path.join(self.work_dir, self.genome + '.extrinsic_hints.gff'))
        return luigi.LocalTarget(hints)

    def run(self):
        logger.info('Beginning GenerateHints Toil pipeline for {}.'.format(self.genome))
        work_dir = os.path.abspath(os.path.join(self.work_dir, 'toil', self.genome))
        toil_options = self.prepare_toil_options(work_dir)
        toil_options.defaultMemory = '8G'  # TODO: don't hardcode this.
        generate_hints(self.genome, self.flat_fasta, self.cfg, self.annotation, self.output().path, toil_options)
        logger.info('Finished GenerateHints Toil pipeline for {}.'.format(self.genome))


class BuildDb(HintsDbTask):
    """
    Constructs the hints database from a series of reduced hints GFFs

    TODO: output() should be way smarter than this. Currently, it only checks if the indices have been created.
    """
    cfg = luigi.Parameter()
    hint_paths = luigi.Parameter()
    flat_fasta_paths = luigi.Parameter()
    genomes = luigi.TupleParameter()
    target_genomes = luigi.TupleParameter()

    def requires(self):
        for genome in self.genomes:
            flat_fasta = self.clone(GenomeFlatFasta, genome=genome, cfg=self.cfg)
            yield flat_fasta
        for genome in self.target_genomes:
            flat_fasta = self.clone(GenomeFlatFasta, genome=genome, cfg=self.cfg)
            annotation = self.cfg['ANNOTATION'].get(genome, None)
            hints = self.clone(GenerateHints, genome=genome, flat_fasta=flat_fasta.output().path,
                               annotation=annotation, cfg=self.cfg)
            yield hints

    def output(self):
        tools.fileOps.ensure_file_dir(self.augustus_hints_db)
        return IndexTarget(self.augustus_hints_db)

    def run(self):
        for genome in self.genomes:
            logger.info('Loading sequence for {} into database.'.format(genome))
            base_cmd = ['load2sqlitedb', '--noIdx', '--clean', '--species={}'.format(genome),
                        '--dbaccess={}'.format(self.augustus_hints_db)]
            tools.procOps.run_proc(base_cmd + [self.flat_fasta_paths[genome]])
            if genome in self.hint_paths:
                logger.info('Loading hints for {} into database.'.format(genome))
                tools.procOps.run_proc(base_cmd + [self.hint_paths[genome]])
        logger.info('Indexing database.')
        cmd = ['load2sqlitedb', '--makeIdx', '--clean', '--dbaccess={}'.format(self.augustus_hints_db)]
        tools.procOps.run_proc(cmd)


class IndexTarget(luigi.Target):
{'fraction_non_alphanumeric': 0.10702154908402473, 'fraction_numerical': 0.0011693301408764408, 'mean_word_length': 2.548004741209008, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1425788', 'n_tokens_mistral': 5576, 'n_tokens_neox': 5266, 'n_words': 1345}