##!/bin/env python
'''farm.py - process data stream on cluster
===========================================


Purpose
-------

This script reads data from stdin and splits it into independent
chunks to be executed on a cluster.

Usage
-----

As a basic, but not very useful example, the following command will
take the input of the file ``go.txt``, split the file by the contents
of the first column and execute a perl command on them::

   cat go.txt | farm.py --split-at-colum=1 perl -p -e "s/GO/gaga/"

Type::

   python farm.py --help

for command line help.

Documentation
-------------

The input on stdin is split for embarrasingly parallel jobs.  The
``--split-at`` options describe how standard input is to be split. A
temporary directory is created in the current directory. This
directory has be visible on the cluster nodes and accessible under the
same name.

The output is written to stdout. Results are returned in the same
order as they are submitted. The script implements a few generic ways
to combine tabular output, for example to avoid duplicating header
lines. The script is also able to handle multiple outputs for jobs.

On error, error messages are echoed and nothing is returned.  The
temporary directory is not deleted to allow manual recovery.

Examples
--------

The following command will split the file "go" at the first column and
execute the command perl -p -e "s/GO/gaga/"::

   cat go | farm.py --split-at-colum=1 perl -p -e "s/GO/gaga/"

The following command will split a fasta file at each entry and
compute an approximate sequence length::

   cat genome.fasta | farm.py --split-at-regex="^>(\S+)" "wc -c"

The following command will split a fasta file at every 10 sequences::

   cat genome.fasta | farm.py --split-at-regex="^>(\S+)" --chunk-size=10 "wc -c"

.. todo::

   implement continuation of jobs
   implement better error messages
   use sge array jobs for job control

Command line options
--------------------

'''

import os
import sys
import re
import glob
import subprocess
import tempfile
import shutil
import stat

from multiprocessing.pool import Pool, ThreadPool

import CGAT.Experiment as E
import CGAT.IOTools as IOTools
import CGAT.Blat as Blat

from CGATPipelines.Pipeline import Cluster as Cluster

try:
    import drmaa
    HAS_DRMAA = True
except (ImportError, RuntimeError):
    HAS_DRMAA = False


def chunk_iterator_lines(infile, args, prefix, use_header=False):
    """split by lines."""

    chunk_size = args[0]
    n = 0
    filename = "%s/%010i.in" % (prefix, n)
    outfile = IOTools.openFile(filename, "w")
    header = None

    for line in infile:
        if line[0] == "#":
            continue

        if not header and n == 0 and use_header:
            header = line
            outfile.write(header)
            continue

        n += 1

        if n % chunk_size == 0:
            outfile.close()
            yield filename
            filename = "%s/%010i.in" % (prefix, n)
            outfile = IOTools.openFile(filename, "w")
            if header:
                outfile.write(header)

        outfile.write(line)
    outfile.close()
    yield filename


def chunk_iterator_column(infile, args, prefix, use_header=False):
    """split at column.

    The table need not be sorted by this column.
    If num_files is given, files will randomly created
    and tags according to column randomly assigned.


    """

    column, max_files = args
    files = IOTools.FilePool()
    header = False

    if max_files:
        map_tag2file = {}

    for line in infile:
        if line[0] == "#":
            continue

        if not header and use_header:
            files.setHeader(line)
            header = True
            continue

        key = line[:-1].split("\t")[column]
        if max_files:
            if key in map_tag2file:
                key = map_tag2file[key]
            else:
                n = "%010i" % (len(map_tag2file) % max_files)
                map_tag2file[key] = n
                key = n

        files.write("%s/%s.in" % (prefix, key), line)

    for filename, count in list(files.items()):
        E.info("created file %s with %i items" % (filename, count))
        yield filename


def chunk_iterator_regex_group(infile, args, prefix, use_header=False):
    """group by regular expression is true.

    Entries need to be consecutive.
    """

    rex = args[0]
    column = args[1]
    chunk_size = args[2]
    last = None
    header = None
    n = chunk_size
    outfile = None
    filename = None

    for line in infile:

        if line[0] == "#":
            continue

        if not header and use_header:
            header = line
            continue

        try:
            this = rex.search(line[:-1]).groups()[0]
        except IndexError:
            if outfile:
                outfile.write(line)
            continue
        except AttributeError:
            if outfile:
                outfile.write(line)
            continue

        if last != this and n >= chunk_size:
            if last:
                outfile.close()
                yield filename

            last = this

            filename = "%s/%s.in" % (prefix, this)
            outfile = IOTools.openFile(filename, "w")
            if header:
                outfile.write(header)
            n = 0

        outfile.write(line)
        n += 1

    if outfile:
        outfile.close()
        yield filename


def chunk_iterator_regex_split(infile, args, prefix, use_header=False):
    """split where regular expression is true.
    """

    rex = args[0]
    chunk_size = args[2]
    max_lines = args[3]

    nlines = 0
    n = 0
    filename = "%s/%010i.in" % (prefix, n)
    outfile = IOTools.openFile(filename, "w")

    for line in infile:

        if line[0] == "#":
            continue
        if rex.search(line[:-1]):
            if n > 0 and (n % chunk_size == 0 or
                          (max_lines and nlines > max_lines)):
                outfile.close()
                yield filename
                filename = "%s/%010i.in" % (prefix, n)
                outfile = IOTools.openFile(filename, "w")
                nlines = 0

            n += 1

        outfile.write(line)
        nlines += 1

    outfile.close()
    yield filename


def chunk_iterator_psl_overlap(infile, args, prefix, use_header=False):
    """iterate over overlapping entries in a psl file."""

    iterator = Blat.BlatIterator(sys.stdin)

    processed_contigs = set()

    merge_distance = args[0]
    last_sbjct_id = None
    sbjct_end = 0
    outfile = None
    filename = None
    while 1:

        match = next(iterator)

        if match is None:
            break

        if match.mSbjctId != last_sbjct_id or \
           match.mSbjctFrom >= (sbjct_end + merge_distance):
            if last_sbjct_id:
                outfile.close()
                yield filename

            if last_sbjct_id != match.mSbjctId and \
               match.mSbjctId in processed_contigs:
                raise ValueError(
                    "input not sorted correctly (contig,start): "
                    "already encountered %s\n%s" %
                    (match.mSbjctId, str(match)))

            last_sbjct_id = match.mSbjctId
            processed_contigs.add(last_sbjct_id)

            sbjct_start = match.mSbjctFrom
            sbjct_end = match.mSbjctTo

        if match.mSbjctFrom < sbjct_start:
            raise ValueError(
                "input not sorted correctly (contig,start): "
                "%i < %i\n%s" %
                (match.mSbjctFrom, sbjct_start, str(match)))

        sbjct_end = max(match.mSbjctTo, sbjct_end)
        outfile.write(str(match) + "\n")

    if outfile:
        outfile.close()
        yield filename


class MapperGlobal:

    def __init__(self, pattern="%06i"):
        self.mMap = {}
        assert "%" in pattern, "please supply a pattern"
        self.mPattern = pattern

    def __call__(self, fn, id):
        if id not in self.mMap:
            self.mMap[id] = self.mPattern % (len(self.mMap) + 1)
        return self.mMap[id]


class MapperLocal:

    def __init__(self, pattern="%06i"):
        self.mMap = {}
        assert "%" in pattern, "please supply a pattern"
        self.mPattern = pattern

    def __call__(self, fn, id):
        key = "%s-%s" % (fn, id)
        if key not in self.mMap:
            self.mMap[key] = self.mPattern % (len(self.mMap) + 1)
        return self.mMap[key]


class MapperEmpty:

    def __init__(self):
        pass

    def __call__(self, fn, id):
        return id


class ResultBuilder:

    """the default result builder for table formatted output.

    field_index :
    """

    def __init__(self,
                 mapper=None,
                 field_index=None,
                 field_name=None,
                 header_regex=None):
        self.mMapper = mapper
        self.mFieldIndex = field_index
        self.mFieldName = field_name
        self.header = None
        self.nfields = None
        self.header_regex = header_regex

    def parseHeader(self, infile, outfile, options):
        """parse header in infile."""
        # skip comments until header
        while 1:
            l = infile.readline()
            if not l:
                break
            if self.header_regex:
                if self.header_regex.search(l):
                    break
            elif l[0] != "#":
                break
            options.stdlog.write(l)

        # print only the first header and check if
        # all the headers are the same.
        if self.header:
            if self.header != l:
                raise ValueError(
                    "inconsistent header in file %s\n"
                    "got=%s\nexpected=%s" % (infile, l, self.header))
        else:
            outfile.write(l)
            self.header = l
            self.nfields = l.count("\t")
            if self.nfields == 0:
                E.warn("only single column in header: %s" % l[:-1])

            if self.mFieldIndex is None and self.mFieldName:
                try:
                    self.mFieldIndex = self.header.split(
                        "\t").index(self.mFieldName)
                except ValueError:
                    E.warn("no mapping, can not find field %s in %s" %
                           (self.mFieldName, self.header))
                    self.mFieldName = None

                E.debug(
                    "substituting field: %s, %s" %
                    (self.mFieldName, self.mFieldIndex))

    def __call__(self, filenames, outfile, options):

        for fi, fn in filenames:
            E.debug("# merging %s" % fn)
            infile = IOTools.openFile(fn, "r")

            if options.output_header:
                self.parseHeader(infile, outfile, options)

            for l in infile:
                nfields = l.count("\t")

                if l[0] == "#":
                    options.stdlog.write(l)
                elif self.nfields is not None and nfields != self.nfields:
                    # validate number of fields in row, raise warning
                    # for those not matching and skip.
                    E.warn(
                        "# line %s has unexpected number of fields: %i != %i" %
                        (l[:-1], nfields, self.nfields))
                else:
                    if self.mFieldIndex is not None:
                        data = l[:-1].split("\t")
                        try:
                            data[self.mFieldIndex] = self.mMapper(
                                fi, data[self.mFieldIndex])
                        except IndexError:
                            raise IndexError(
                                "can not find field %i in %s" %
                                (self.mFieldIndex, l))
                        l = "\t".join(data) + "\n"

                    outfile.write(l)
            infile.close()


class ResultBuilderPSL(ResultBuilder):

    """Result builder for psl tables. Here, column 9,
    the query id, is substituted."""

    def __init__(self, *args, **kwargs):
        ResultBuilder.__init__(self, *args, **kwargs)
        self.mFieldIndex = 9
        self.mFirst = True

    def parseHeader(self, infile, outfile, options):
        """parse header in infile."""
        # skip comments until header
        while 1:
            l = infile.readline()
            if not l or l[0] != "#":
                break
            options.stdlog.write(l)

        if l.startswith("psLayout version 3"):
            if self.mFirst:
                outfile.write(l)
                for x in range(0, 4):
                    l = infile.readline()
                    outfile.write(l)
                self.mFirst = False
            else:
                for x in range(0, 4):
                    l = infile.readline()


class ResultBuilderFasta(ResultBuilder):

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

    def __call__(self, filenames, outfile, options):
        for fi, fn in filenames:
            infile = IOTools.openFile(fn, "r")
            for l in infile:
                if l[0] == "#":
                    options.stdlog.write(l)
                    continue
                elif l[0] == ">":
                    x = re.search(">(\S+)", l[:-1])
                    id = self.mMapper(fi, x.groups()[0])
                    l = ">%s%s" % (id, l[x.end(0):])
                outfile.write(l)
            infile.close()


class ResultBuilderBinary(ResultBuilder):

    '''simply concatenate output files (without any parsing).'''

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

    def __call__(self, filenames, outfile, options):
        for fi, fn in filenames:
            shutil.copyfileobj(IOTools.openFile(fn, "r"), outfile)


class ResultBuilderCopies(ResultBuilder):

    '''create indexed copiers.'''

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

    def __call__(self, filenames, outfile, options):
        idx = 0
        base, ext = os.path.splitext(outfile.name)
        for fi, fn in filenames:
            idx += 1
            shutil.copyfile(fn, base + ".%i" % idx + ext)


class ResultBuilderLog(ResultBuilder):

    """processor for log files."""

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

    def __call__(self, filenames, outfile, options):
        for fi, fn in filenames:
            infile = IOTools.openFile(fn, "r")
            outfile.write(
                "######### logging output for %s ###################\n" % fi)
            for l in infile:
                outfile.write(l)
            infile.close()