#!/usr/bin/env python
# -*- coding: UTF-8 -*-

import os
import os.path as op
import math
import sys
import logging

from itertools import groupby, islice, cycle, izip

from Bio import SeqIO
from jcvi.apps.base import OptionParser, ActionDispatcher, sh, debug, need_update, \
            mkdir, popen
debug()


FastaExt = ("fasta", "fa", "fna", "cds", "pep", "faa", "fsa", "seq", "nt", "aa")
FastqExt = ("fastq", "fq")


class BaseFile (object):

    def __init__(self, filename):

        self.filename = filename
        if filename:
            logging.debug("Load file `{0}`".format(filename))


class LineFile (BaseFile, list):
    """
    Generic file parser for line-based files
    """
    def __init__(self, filename, comment=None, load=False):

        super(LineFile, self).__init__(filename)

        if load:
            fp = must_open(filename)
            self.lines = [l.strip() for l in fp if l[0]!=comment]
            logging.debug("Load {0} lines from `{1}`.".\
                        format(len(self.lines), filename))


class DictFile (BaseFile, dict):
    """
    Generic file parser for multi-column files, keyed by a particular index.
    """
    def __init__(self, filename, keypos=0, valuepos=1, delimiter=None,
                       strict=True, cast=None):

        super(DictFile, self).__init__(filename)

        fp = must_open(filename)
        ncols = max(keypos, valuepos) + 1
        thiscols = 0
        for lineno, row in enumerate(fp):
            row = row.rstrip()
            atoms = row.split(delimiter)
            thiscols = len(atoms)
            if thiscols < ncols:
                action = "Aborted" if strict else "Skipped"

                msg = "Must contain >= {0} columns.  {1}.\n".format(ncols, action)
                msg += "  --> Line {0}: {1}".format(lineno + 1, row)
                logging.error(msg)
                if strict:
                    sys.exit(1)
                else:
                    continue

            key = atoms[keypos]
            value = atoms[valuepos] if (valuepos is not None) else atoms
            if cast:
                value = cast(value)
            self[key] = value

        assert thiscols, "File empty"
        self.ncols = thiscols
        logging.debug("Imported {0} records from `{1}`.".\
                    format(len(self), filename))


class SetFile (BaseFile, set):

    def __init__(self, filename, column=-1, delimiter=None):
        super(SetFile, self).__init__(filename)
        fp = open(filename)
        for row in fp:
            if not row.strip():
                continue
            keys = [x.strip() for x in row.split(delimiter)]
            if column >= 0:
                keys = [keys[column]]
            self.update(keys)


class FileShredder (object):
    """
    Same as rm -f *
    """
    def __init__(self, filelist):

        filelist = [x for x in filelist if x and op.exists(x)]
        cmd = "rm -rf {0}".format(" ".join(filelist))
        sh(cmd)


class FileMerger (object):
    """
    Same as cat * > filename
    """
    def __init__(self, filelist, outfile):

        self.filelist = filelist
        self.outfile = outfile
        self.ingz = filelist[0].endswith(".gz")
        self.outgz = outfile.endswith(".gz")

    def merge(self, checkexists=False):
        outfile = self.outfile
        if checkexists and not need_update(self.filelist, outfile):
            logging.debug("File `{0}` exists. Merge skipped.".format(outfile))
            return

        files = " ".join(self.filelist)
        ingz, outgz = self.ingz, self.outgz
        if ingz and outgz:  # can merge gz files directly
            cmd = "cat {0} > {1}".format(files, outfile)
            sh(cmd)
        else:
            cmd = "zcat" if self.ingz else "cat"
            cmd += " " + files
            sh(cmd, outfile=outfile)

        return outfile


class FileSplitter (object):

    def __init__(self, filename, outputdir=None, format="fasta", mode="cycle"):
        self.filename = filename
        self.outputdir = outputdir
        self.mode = mode

        self.format = format
        logging.debug("format is %s" % format)

        guessedformat = self._guess_format(filename)
        if format != guessedformat:
            logging.warn("warn: format guessed from suffix - {0}"\
                          .format(guessedformat))

        if format in ("fasta", "fastq"):
            self.klass = "seqio"
        else:
            self.klass = "txt"

        mkdir(outputdir)

    def _open(self, filename):

        if self.klass == "seqio":
            handle = SeqIO.parse(open(filename), self.format)
        else:
            handle = open(filename)
        return handle

    @property
    def num_records(self):
        handle = self._open(self.filename)
        return sum(1 for x in handle)

    def _guess_format(self, filename):
        root, ext = op.splitext(filename)
        ext = ext.strip(".")

        if ext in FastaExt:
            format = "fasta"
        elif ext in FastqExt:
            format = "fastq"
        else:
            format = "txt"
        return format

    def _batch_iterator(self, N=1):
        """Returns N lists of records.

        This can be used on any iterator, for example to batch up
        SeqRecord objects from Bio.SeqIO.parse(...), or to batch
        Alignment objects from Bio.AlignIO.parse(...), or simply
        lines from a file handle.

        This is a generator function, and it returns lists of the
        entries from the supplied iterator.  Each list will have
        batch_size entries, although the final list may be shorter.
        """
        batch_size = math.ceil(self.num_records / float(N))
        handle = self._open(self.filename)
        while True:
            batch = list(islice(handle, batch_size))
            if not batch:
                break
            yield batch

    @classmethod
    def get_names(cls, filename, N):
        root, ext = op.splitext(op.basename(filename))

        names = []
        pad0 = len(str(int(N - 1)))
        for i in xrange(N):
            name = "{0}_{1:0{2}d}{3}".format(root, i, pad0, ext)
            names.append(name)

        return names

    def write(self, fw, batch):
        if self.klass == "seqio":
            count = SeqIO.write(batch, fw, self.format)
        else:
            for line in batch:
                fw.write(line)
            count = len(batch)
        return count

    def split(self, N, force=False):
        """
        There are two modes of splitting the records
        - batch: splitting is sequentially to records/N chunks
        - cycle: placing each record in the splitted files and cycles

        use `cycle` if the len of the record is not evenly distributed
        """
        mode = self.mode
        assert mode in ("batch", "cycle", "optimal")
        logging.debug("set split mode=%s" % mode)

        self.names = self.__class__.get_names(self.filename, N)
        if self.outputdir:
            self.names = [op.join(self.outputdir, x) for x in self.names]

        if not need_update(self.filename, self.names) and not force:
            logging.error("file %s already existed, skip file splitting" % \
                    self.names[0])
            return

        filehandles = [open(x, "w") for x in self.names]

        if mode == "batch":
            for batch, fw in zip(self._batch_iterator(N), filehandles):
                count = self.write(fw, batch)
                logging.debug("write %d records to %s" % (count, fw.name))

        elif mode == "cycle":
            handle = self._open(self.filename)
            for record, fw in izip(handle, cycle(filehandles)):
                count = self.write(fw, [record])

        elif mode == "optimal":
            """
            This mode is based on Longest Processing Time (LPT) algorithm:

            A simple, often-used algorithm is the LPT algorithm (Longest
            Processing Time) which sorts the jobs by its processing time and
            then assigns them to the machine with the earliest end time so far.
            This algorithm achieves an upper bound of 4/3 - 1/(3m) OPT.

            Citation: <http://en.wikipedia.org/wiki/Multiprocessor_scheduling>
            """
            endtime = [0] * N
            handle = self._open(self.filename)
            for record in handle:
                mt, mi = min((x, i) for (i, x) in enumerate(endtime))
                fw = filehandles[mi]
                count = self.write(fw, [record])
                endtime[mi] += len(record)

        for fw in filehandles:
            fw.close()


def longest_unique_prefix(query, targets, remove_self=True):
    """
    Find the longest unique prefix for filename, when compared against a list of
    filenames. Useful to simplify file names in a pool of files. See usage in
    formats.fasta.pool().
    """
    query = op.basename(query)
    targets = [op.basename(x) for x in targets]
    prefix_lengths = [len(op.commonprefix([query, name])) for name in targets]
    if remove_self and len(query) in prefix_lengths:
        prefix_lengths.remove(len(query))
    longest_length = max(prefix_lengths)
    return query[:longest_length + 1]


def check_exists(filename, oappend=False):
    """
    Avoid overwriting some files accidentally.
    """
    if op.exists(filename):
        if oappend:
            return oappend
        logging.error("`{0}` found, overwrite (Y/N)?".format(filename))
        overwrite = (raw_input() == 'Y')
    else:
        overwrite = True

    return overwrite


def must_open(filename, mode="r", checkexists=False, skipcheck=False, \
            oappend=False):
    """
    Accepts filename and returns filehandle.

    Checks on multiple files, stdin/stdout/stderr, .gz or .bz2 file.
    """
    if isinstance(filename, list):
        assert "r" in mode

        if filename[0].endswith(".gz") or filename[0].endswith(".bz2"):
            filename = " ".join(filename)  # allow opening multiple gz/bz2 files
        else:
            import fileinput
            return fileinput.input(filename)

    if filename in ("-", "stdin"):
        assert "r" in mode
        fp = sys.stdin

    elif filename == "stdout":
        assert "w" in mode
        fp = sys.stdout

    elif filename == "stderr":
        assert "w" in mode
        fp = sys.stderr

    elif filename == "tmp" and mode == "w":
        from tempfile import NamedTemporaryFile
        fp = NamedTemporaryFile(delete=False)

    elif filename.endswith(".gz"):
        if 'r' in mode:
            cmd = "zcat {0}".format(filename)
            fp = popen(cmd, debug=False)
        elif 'w' in mode:
            import gzip
            fp = gzip.open(filename, mode)

    elif filename.endswith(".bz2"):
        if 'r' in mode:
            cmd = "bzcat {0}".format(filename)
            fp = popen(cmd, debug=False)
        elif 'w' in mode:
            import bz2
            fp = bz2.BZ2File(filename, mode)

    else:
        if checkexists:
            assert mode == "w"
            overwrite = (not op.exists(filename)) if skipcheck \
                        else check_exists(filename, oappend)
            if overwrite:
                if oappend:
                    fp = open(filename, "a")
                else:
                    fp = open(filename, "w")
            else:
                logging.debug("File `{0}` already exists. Skipped."\
                        .format(filename))
                return None
        else:
            fp = open(filename, mode)

    return fp


def write_file(filename, contents, meta="file", skipcheck=False, tee=False):
    meta_choices = ("file", "run script", "python script")
    assert meta in meta_choices, "meta must be one of {0}".\
                    format("|".join(meta_choices))

    if meta == "run script":
        contents = "#!/bin/bash\n\n" + contents.strip()
    elif meta == "python script":
        contents = "#!/usr/bin/env python\n# -*- coding: UTF-8 -*-\n\n" + \
                    contents.strip()

    fw = must_open(filename, "w", checkexists=True, skipcheck=skipcheck)
    if fw:
        print >> fw, contents
        fw.close()
    if tee:
        print >> sys.stderr, contents

    message = "{0} written to `{1}`.".format(meta, filename)
    logging.debug(message.capitalize())
    if meta == "run script":
        sh("chmod u+x {0}".format(filename))


def read_until(handle, start):
    # read each line until a certain start, then puts the start tag back
    while 1:
        pos = handle.tell()
        line = handle.readline()
        if not line:
            break
        if line.startswith(start):
            handle.seek(pos)
            return
    #raise EOFError, "%s tag cannot be found"


def read_block(handle, signal):
    """
    Useful for reading block-like file formats, for example FASTA or OBO file,
    such file usually startswith some signal, and in-between the signals are a
    record
    """
    signal_len = len(signal)
    it = (x[1] for x in groupby(handle,
        key=lambda row: row.strip()[:signal_len] == signal))
    found_signal = False
    for header in it:
        header = header.next().strip()
        if header[:signal_len] != signal:
            continue
        found_signal = True
        seq = list(s.strip() for s in it.next())
        yield header, seq

    if not found_signal:
        handle.seek(0)
        seq = list(s.strip() for s in handle)
        yield None, seq


def is_number(s, cast=float):
    """
    Check if a string is a number. Use cast=int to check if s is an integer.
    """
    try:
        cast(s) # for int, long and float
    except ValueError:
        return False

    return True


def get_number(s, cast=int):
    """
    Try to get a number out of a string, and cast it.
    """
    import string
    d = "".join(x for x in str(s) if x in string.digits)
    return cast(d)


def flexible_cast(s):
    if is_number(s, cast=int):
        return int(s)
    elif is_number(s, cast=float):
        return float(s)
    return s


def main():

    actions = (
        ('split', 'split large file into N chunks'),
        ('reorder', 'reorder columns in tab-delimited files'),
        ('flatten', 'convert a list of IDs into one per line'),
        ('group', 'group elements in a table based on key (groupby) column'),
        ('setop', 'set operations on files'),
        ('join', 'join tabular-like files based on common column'),
        ('subset', 'subset tabular-like files based on common column'),
        ('truncate', 'remove lines from end of file'),
        ('append', 'append a column with fixed value'),
            )
    p = ActionDispatcher(actions)
    p.dispatch(globals())


def append(args):
    """
    %prog append csvfile [tag]

    Append a column with fixed value. If tag is missing then just append the
    filename.
    """
    p = OptionParser(append.__doc__)
    p.set_sep()
    p.set_outfile()
    opts, args = p.parse_args(args)

    nargs = len(args)
    if nargs not in (1, 2):
        sys.exit(not p.print_help())

    csvfile = args[0]
    tag = args[1] if nargs == 2 else csvfile
    fp = must_open(csvfile)
    fw = must_open(opts.outfile, "w")
    for row in fp:
        row = row.rstrip("\r\n")
        row = opts.sep.join((row, tag))
        print >> fw, row


def truncate(args):
    """
    %prog truncate linecount filename

    Remove linecount lines from the end of the file in-place. Borrowed from:
    <http://superuser.com/questions/127786/how-to-remove-the-last-2-lines-of-a-very-large-file>
    """
    p = OptionParser(truncate.__doc__)
    opts, args = p.parse_args(args)

    if len(args) != 2:
        sys.exit(not p.print_help())

    number, filename = args
    number = int(number)
    count = 0

    f = open(filename, "r+b")
    f.seek(0, os.SEEK_END)
    while f.tell() > 0:
        f.seek(-1, os.SEEK_CUR)
        char = f.read(1)
        if char == '\n':
            count += 1
        if count == number + 1:
            f.truncate()
            print >> sys.stderr, "Removed {0} lines from end of file".format(number)
            return number

        f.seek(-1, os.SEEK_CUR)

    if count < number + 1:
        print >> sys.stderr, "No change: requested removal would leave empty file"
        return -1


def flatten(args):
    """
    %prog flatten filename > ids

    Convert a list of IDs (say, multiple IDs per line) and move them into one
    per line.

    For example, convert this, to this:
    A,B,C                    | A
    1                        | B
    a,4                      | C
                             | 1
                             | a
                             | 4

    If multi-column file with multiple elements per column, zip then flatten like so:
    A,B,C    2,10,gg         | A,2
    1,3      4               | B,10
                             | C,gg
                             | 1,4
                             | 3,na
    """
    from itertools import izip_longest

    p = OptionParser(flatten.__doc__)
    p.set_sep(sep=",")
    p.add_option("--zipflatten", default=None, dest="zipsep",
                 help="Specify if columns of the file should be zipped before" +
                 " flattening. If so, specify delimiter separating column elements" +
                 " [default: %default]")
    opts, args = p.parse_args(args)

    if len(args) != 1:
        sys.exit(not p.print_help())

    tabfile, = args
    zipsep = opts.zipsep

    fp = must_open(tabfile)
    for row in fp:
        if zipsep:
            row = row.rstrip()
            atoms = row.split(opts.sep)
            frows = []
            for atom in atoms:
                frows.append(atom.split(zipsep))
            print "\n".join([zipsep.join(x) for x in list(izip_longest(*frows, fillvalue="na"))])
        else:
            print row.strip().replace(opts.sep, "\n")


def group(args):
    """
    %prog group tabfile > tabfile.grouped

    Given a tab-delimited file, either group all elements within the file or
    group the elements in the value column(s) based on the key (groupby) column

    For example, convert this | into this
    ---------------------------------------
    a	2    3    4           | a,2,3,4,5,6
    a	5    6                | b,7,8
    b	7    8                | c,9,10,11
    c	9                     |
    c 	10   11               |

    If grouping by a particular column,
    convert this              | into this:
    ---------------------------------------------
    a	2    3    4           | a	2,5   3,6   4
    a	5    6                | b	7     8
    b	7    8                | c	9,10  11
    c	9                     |
    c 	10   11               |

    By default, it uniqifies all the grouped elements
    """
    from jcvi.utils.cbook import AutoVivification
    from jcvi.utils.grouper import Grouper

    p = OptionParser(group.__doc__)
    p.set_sep()
    p.add_option("--groupby", default=None, type='int',
                 help="Default column to groupby [default: %default]")
    p.add_option("--groupsep", default=',',
                 help="Separator to join the grouped elements [default: `%default`]")
    p.add_option("--nouniq", default=False, action="store_true",
                 help="Do not uniqify the grouped elements [default: %default]")
    opts, args = p.parse_args(args)

    if len(args) != 1:
        sys.exit(not p.print_help())

    tabfile, = args
    sep = opts.sep
    groupby = opts.groupby
    groupsep = opts.groupsep

    cols = []
    grouper = AutoVivification() if groupby is not None else Grouper()
    fp = must_open(tabfile)
    for row in fp:
        row = row.rstrip()
        atoms = row.split(sep)
        if groupby is not None:
            if len(cols) < len(atoms):
                cols = [x for x in xrange(len(atoms))]
            if groupby not in cols:
                logging.error("groupby col index `{0}` is out of range".format(groupby))
                sys.exit()

            key = atoms[groupby]
            for col in cols:
                if col == groupby:
                    continue
                if not grouper[key][col]:
                    grouper[key][col] = [] if opts.nouniq else set()
                if col < len(atoms):
                    if groupsep in atoms[col]:
                        for atom in atoms[col].split(groupsep):
                            if opts.nouniq:
                                grouper[key][col].append(atom)
                            else:
                                grouper[key][col].add(atom)
                    else:
                        if opts.nouniq:
                            grouper[key][col].append(atoms[col])
                        else:
                            grouper[key][col].add(atoms[col])
        else:
            grouper.join(*atoms)

    for key in grouper:
        if groupby is not None:
            line = []
            for col in cols:
                if col == groupby:
                    line.append(key)
                elif col in grouper[key].keys():
                    line.append(groupsep.join(grouper[key][col]))
                else:
                    line.append("na")
            print sep.join(line)
        else:
            print groupsep.join(key)


def reorder(args):
    """
    %prog reorder tabfile 1,2,4,3 > newtabfile

    Reorder columns in tab-delimited files. The above syntax will print out a
    new file with col-1,2,4,3 from the old file.
    """
    import csv

    p = OptionParser(reorder.__doc__)
    p.set_sep()
    opts, args = p.parse_args(args)

    if len(args) != 2:
        sys.exit(not p.print_help())

    tabfile, order = args
    sep = opts.sep
    order = [int(x) - 1 for x in order.split(",")]
    reader = csv.reader(must_open(tabfile), delimiter=sep)
    writer = csv.writer(sys.stdout, delimiter=sep)
    for row in reader:
        newrow = [row[x] for x in order]
        writer.writerow(newrow)


def split(args):
    """
    %prog split file outdir N

    Split file into N records. This allows splitting FASTA/FASTQ/TXT file
    properly at boundary of records. Split is useful for parallelization
    on input chunks.

    Option --mode is useful on how to break into chunks.
    1. chunk - chunk records sequentially, 1-100 in file 1, 101-200 in file 2, etc.
    2. cycle - chunk records in Round Robin fashion
    3. optimal - try to make split file of roughly similar sizes, using LPT
    algorithm. This is the default.
    """
    p = OptionParser(split.__doc__)
    mode_choices = ("batch", "cycle", "optimal")
    p.add_option("--all", default=False, action="store_true",
            help="split all records [default: %default]")
    p.add_option("--mode", default="optimal", choices=mode_choices,
            help="Mode when splitting records [default: %default]")
    p.add_option("--format", default="fasta", choices=("fasta", "fastq", "txt"),
            help="input file format [default: %default]")

    opts, args = p.parse_args(args)

    if len(args) != 3:
        sys.exit(not p.print_help())

    mode = opts.mode
    filename, outdir, N = args
    fs = FileSplitter(filename, outputdir=outdir, format=opts.format, mode=mode)

    if opts.all:
        logging.debug("option -all override N")
        N = fs.num_records
    else:
        N = int(N)
        assert N > 0, "N must be > 0"

    logging.debug("split file into %d chunks" % N)
    fs.split(N)

    return fs


def join(args):
    """
    %prog join file1.txt(pivotfile) file2.txt ..

    Join tabular-like files based on common column.
    --column specifies the column index to pivot on.
      Use comma to separate multiple values if the pivot column is different
      in each file. Maintain the order in the first file.
    --sep specifies the column separators, default to tab.
      Use comma to separate multiple values if the column separator is different
      in each file.
    """
    from jcvi.utils.iter import flatten

    p = OptionParser(join.__doc__)
    p.add_option("--column", default="0",
                 help="0-based column id, multiple values allowed [default: %default]")
    p.set_sep(multiple=True)
    p.add_option("--noheader", default=False, action="store_true",
                 help="Do not print header [default: %default]")
    p.add_option("--na", default="na",
                 help="Value for unjoined data [default: %default]")
    p.add_option("--keysep", default=",",
                 help="specify separator joining multiple elements in the key column"
                 + " of the pivot file [default: %default]")
    p.set_outfile()

    opts, args = p.parse_args(args)
    nargs = len(args)

    keysep = opts.keysep

    if len(args) < 2:
        sys.exit(not p.print_help())

    na = opts.na
    c = opts.column
    if "," in c:
        cc = [int(x) for x in c.split(",")]
    else:
        cc = [int(c)] * nargs

    assert len(cc) == nargs, "Column index number != File number"

    s = opts.sep
    if "," in s:
        ss = [x for x in s.split(",")]
    else:
        ss = [s] * nargs

    assert len(ss) == nargs, "column separator number != File number"

    # Maintain the first file line order, and combine other files into it
    pivotfile = args[0]
    files = [DictFile(f, keypos=c, valuepos=None, delimiter=s) \
                        for f, c, s in zip(args, cc, ss)]
    otherfiles = files[1:]
    header = "\t".join(flatten([op.basename(x.filename)] * x.ncols \
                        for x in files))

    fp = must_open(pivotfile)
    fw = must_open(opts.outfile, "w")
    if not opts.noheader:
        print >> fw, header

    for row in fp:
        row = row.rstrip()
        atoms = row.split(ss[0])
        newrow = atoms
        key = atoms[cc[0]]
        keys = key.split(keysep) if keysep in key else [key]
        for d in otherfiles:
            drows = list()
            for key in keys:
                drows.append(d.get(key, [na] * d.ncols))
            drow = [keysep.join(x) for x in list(zip(*drows))]
            newrow += drow
        print >> fw, "\t".join(newrow)


def subset(args):
    """
    %prog subset file1.txt(pivotfile) file2.txt ..

    subset tabular-like file1 based on common column with file 2.
    Normally file1 should have unique row entries.
    If more than one file2 are provided, they must have same column separators.
    Multiple file2's will be concatenated in the output.

    --column specifies the column index (0-based) to pivot on.
      Use comma to separate multiple values if the pivot column is different
      in each file. Maintain the order in the first file.
    --sep specifies the column separators, default to tab.
      Use comma to separate multiple values if the column separator is different
      in each file.
    """

    p = OptionParser(subset.__doc__)
    p.add_option("--column", default="0",
                 help="0-based column id, multiple values allowed [default: %default]")
    p.set_sep(multiple=True)
    p.add_option("--pivot", default=1, type="int",
                 help="1 for using order in file1, 2 for using order in \
                    file2 [default: %default]")
    p.set_outfile()

    opts, args = p.parse_args(args)
    nargs = len(args)

    if len(args) < 2:
        sys.exit(not p.print_help())

    c = opts.column
    if "," in c:
        cc = [int(x) for x in c.split(",")]
        assert len(set(cc[1:])) == 1, \
            "Multiple file2's must have same column index."
        cc = cc[0:2]
    else:
        cc = [int(c)] * 2

    s = opts.sep
    if "," in s:
        ss = [x for x in s.split(",")]
        assert len(set(cc[1:])) == 1, \
            "Multiple file2's must have same column separator."
        ss = ss[0:2]
    else:
        ss = [s] * 2

    if nargs>2:
        file2 = FileMerger(args[1:], outfile="concatenatedFile2").merge()
    else:
        file2 = args[1]
    newargs = [args[0], file2]

    files = [DictFile(f, keypos=c, valuepos=None, delimiter=s) \
                        for f, c, s in zip(newargs, cc, ss)]

    pivot = 0 if opts.pivot==1 else 1
    fp = open(newargs[pivot])
    fw = must_open(opts.outfile, "w")

    for row in fp:
        row = row.rstrip()
        atoms = row.split(ss[pivot])
        key = atoms[cc[pivot]]
        d = files[1-pivot]
        if key in d:
            print >> fw, ss[0].join(files[0][key])

    if nargs>2:
        FileShredder([file2])


def setop(args):
    """
    %prog setop "fileA & fileB" > newfile

    Perform set operations, except on files. The files (fileA and fileB) contain
    list of ids. The operator is one of the four:

    |: union (elements found in either file)
    &: intersection (elements found in both)
    -: difference (elements in fileA but not in fileB)
    ^: symmetric difference (elementes found in either set but not both)

    Please quote the argument to avoid shell interpreting | and &.
    """
    p = OptionParser(setop.__doc__)
    p.add_option("--column", default=0, type="int",
                 help="The column to extract, 0-based, -1 to disable [default: %default]")
    opts, args = p.parse_args(args)

    if len(args) != 1:
        sys.exit(not p.print_help())

    statement, = args
    fa, op, fb = statement.split()
    assert op in ('|', '&', '-', '^')

    column = opts.column
    fa = SetFile(fa, column=column)
    fb = SetFile(fb, column=column)

    if op == '|':
        t = fa | fb
    elif op == '&':
        t = fa & fb
    elif op == '-':
        t = fa - fb
    elif op == '^':
        t = fa ^ fb

    for x in sorted(t):
        print x


if __name__ == '__main__':
    main()
