#!/usr/bin/env python
""" Make a ds9 region file from a SExtractor catalogue."""
from __future__ import division, print_function, unicode_literals
try:
    unicode
except NameError:
    import pickle
    unicode = basestring = str
    xrange = range
else:
    import cPickle as pickle
import numpy as np

usage = """\
Usage: sex2DS9reg sextractor_catalogue_filename [DS9region_filename]
Create a DS9 region file from a SExtractor output catalogue. It must
contain either X_IMAGE and Y_IMAGE, or XWIN_IMAGE and YWIN_IMAGE. If
A, B and THETA columns are also present, then ellipse regions are
drawn, otherwise just points.
"""


def readtxt(fh,
            sep=None,
            usecols=None,
            comment='#',
            skip=0,
            arrays=True,
            names=None,
            readnames=False,
            converters=None,
            mintype=int):
    """ Reads columns from a text file into arrays, converting to int,
    float or str where appropriate.
    By default the column separator is whitespace. `rows` can be
    either an input filename or an iterable (e.g. a file object, list
    or iterator).
    Parameters
    ----------
    rows : filename or iterable object
        Input data.
    sep : str (default `None`)
        A string used to separate items on a row (also known as a
        delimiter). Default is None, which means whitespace.
    usecols : int or tuple of ints, optional
        Indices of columns to be read. By default all columns are read.
    comment : str (default `#`)
        Character marking the start of a comment.
    skip : int (default `0`)
        Number of rows to skip (not counting commented or blank lines)
        before reading data.
    arrays : bool (`True`)
        If True, all columns are converted to Numpy arrays.  If False,
        columns are returned as lists.
    names : str or sequence of str (default `None`)
        If `names` is given and `arrays` is True, the data are placed
        in a Numpy record array with field names given by `names`. Can
        also be a single string of comma-separated values.
    readnames : bool (`False`)
        If `readnames` is True the first line of the file is read to
        find the field names. This overrides the `names` keyword.
    converters : dict (`None`)
        Functions to apply to each entry of a column. Each (key,value)
        pair gives the column index (key) and the function to be
        applied to each entry in that column (value).
    Returns either structured array or lists.
    Examples
    --------
    >>> list_of_all_cols = readtxt('filename')
    >>> ninthcol, fifthcol = readtxt('filename', sep=',', usecols=(8,4)])
    >>> firstcol = readtxt('filename', comment='%', usecols=[0])
    >>> recarray = readtxt('filename', sep=',', usecols=(1,3), names='x,y'])
    """
    if mintype == float:
        typedict = {float: lambda x: str(x).strip()}
    elif mintype == int:
        typedict = {int: float, float: lambda x: str(x).strip()}
    else:
        raise ValueError('Unknown minimum type %s' % mintype)

    def convert(row, funcs):
        # convert each item in a row to int, float or str.
        for i, item in enumerate(row):
            while True:
                try:
                    row[i] = funcs[i](item)
                except ValueError:
                    # update the list of converters
                    try:
                        funcs[i] = typedict[funcs[i]]
                    except KeyError:
                        raise ValueError('Converter %s failed '
                                         'on %r' % (funcs[i], item))
                else:
                    break
        return row, funcs

    needclose = False
    if isinstance(fh, basestring):
        if fh.endswith('.gz'):
            import gzip
            fh = gzip.open(fh)
        else:
            fh = open(fh)
        needclose = True

    data = iter(fh)

    if comment is not None:
        len_comment = len(comment)

    if names and isinstance(names, basestring):
        names = [n.strip() for n in str(names).split(',')]
    elif names:
        names = map(str, names)

    skipped = 0
    out = []
    # main loop to read data
    for irow, row in enumerate(data):
        if comment is not None:
            row = row.split(comment)[0]
        row = row.lstrip()
        if not row: continue
        if skipped < skip:
            skipped += 1
            continue
        row = row.split(sep)
        if readnames:
            names = [str(r.strip()) for r in row]
            readnames = False
            continue
        if not out:
            # first row with data, so initialise converters
            funcs = [mintype] * len(row)
            if converters is not None:
                for i in converters:
                    funcs[i] = converters[i]
            if usecols is not None:
                funcs = [funcs[i] for i in usecols]
        if usecols is not None:
            try:
                row = [row[i] for i in usecols]
            except IndexError:
                raise IndexError('Columns indices: %s, but only %i entries in '
                                 'this row!' % (usecols, len(row)))
        try:
            row, funcs = convert(row, funcs)
        except IndexError:
            # Probably there are more items in this row than in
            # previous rows. This usually indicates a problem, so
            # raise an error.
            raise IndexError('Too many items on row %i: %s' % (irow + 1, row))

        if names:
            assert len(row) == len(names), '%i, %i, %s ' % (len(names),
                                                            irow + 1, row)
        out.append(row)

    if needclose:
        fh.close()

    # rows to columns, truncating to number of words on shortest line.
    if arrays:
        if names is not None:
            out = np.rec.fromrecords(out, names=names)
        else:
            out = [np.array(c) for c in zip(*out)]
    else:
        out = [list(c) for c in zip(*out)]

    if len(out) == 1 and names is None:
        return out[0]
    else:
        return out


def readsex(filename, catnum=None):
    """ Read a sextractor catalogue into a Numpy record array.
    Parameters
    ----------
    filename : str
      Sextractor output catalogue name
    catnum : int, optional
      If the Sextractor file is in LDAC_FITS format and contains more
      than one catalogue, this option specifies the catalogue number.
    Returns
    -------
    s : numpy record array
      Record array with field names the same as those in the
      sextractor catalogue.
    """
    fh = open(filename)
    # get the header
    row = fh.next()
    while not row.strip():
        row = fh.next()
    if row[8] == '=':
        fh.close()
        # assume a fits file
        try:
            import pyfits
        except ImportError:
            import astropy.io.fits as pyfits
        fh = pyfits.open(filename)
        if len(fh) > 3 and catnum is None:
            raise ValueError("specify catalogue number")
        elif catnum is not None:
            return pyfits.getdata(filename, catnum * 2).view(np.recarray)
        else:
            return pyfits.getdata(filename, 2).view(np.recarray)
    hd = []