from .StarList import StarList
from .file_helpers import *
from .daofiles import parse_dao_hdr, write_dao_header, DAO_file_firstline, DAO
from .file_helpers import as_starlist
import pandas as pd
import re

_ds9_regexp = re.compile(
    r'[+-]? *circle[( ] *([+-]?\d+[.]?\d*) *[, ] *([+-]?\d+[.]?\d*).+#.*id *= *(\d+)')
_ds9_wcs_regexp = re.compile(
    r'[+-]? *circle[( ] *([+-]?\d+:\d+:\d+[.]?\d*) *[, ] *([+-]?\d+:\d+:\d+[.]?\d*).+#.*id *= *(\d+)')
_ds9_no_id_regexp = re.compile(
    r'[+-]? *circle[( ] *([+-]?\d+[.]?\d*) *[, ] *([+-]?\d+[.]?\d*)')
_ds9_no_id_wcs_regexp = re.compile(
    r'[+-]? *circle[( ] *([+-]?\d+:\d+:\d+[.]?\d*) *[, ] *([+-]?\d+:\d+:\d+[.]?\d*)')
_ds9_system_wcs = re.compile('fk4|fk5|J2000|B1950|ICRS', re.IGNORECASE)
_ds9_system_xy = re.compile('PHYSICAL|IMAGE', re.IGNORECASE)


def read_ds9_regions(file):
    # type: (object) -> StarList
    """
    Reads ds9 region
    :param file: filename or open input stream
    :return: StarList object

    Returned object has columns id, x, y, auto_id

    Boolean column auto_id indicates weather id for item is read from file (#id=xxx comment) or
    generated by function.
    """

    f, to_close = get_stream(file, 'rt')
    # s = StarList.new()
    data = []
    data_noid = []
    dao_hdr1 = None
    hdr = None
    sys_wcs, sys_xy = (1,2)
    system = None
    for line in f:
        if line[0] == '#':
            if line[1:11] == DAO_file_firstline[:10]:  # dao header found in comment
                dao_hdr1 = line
                continue
            if dao_hdr1 is not None:  # second line of dao header
                hdr = parse_dao_hdr(dao_hdr1, line, '#')
        else:
            if system is None:
                if _ds9_system_wcs.search(line):
                    system = sys_wcs
                elif _ds9_system_xy.search(line):
                    system = sys_xy
                pass
            m = _ds9_regexp.search(line)
            if m is not None:  # s[id] = (id, x, y)
                data.append([int(m.group(3)), float(m.group(1)), float(m.group(2))])
            else:
                m = _ds9_wcs_regexp.search(line)
                if m is not None:  # s[id] = (id, ra, dec)
                    data.append([int(m.group(3)), str(m.group(1)), str(m.group(2))])
                else:
                    m = _ds9_no_id_regexp.search(line)  # s[?] = (x, y)
                    if m is not None:
                        data_noid.append([float(m.group(1)), float(m.group(2))])
                    else:
                        m = _ds9_no_id_wcs_regexp.search(line)  # s[?] = (ra, dec)
                        if m is not None:
                            data_noid.append([str(m.group(1)), str(m.group(2))])

        dao_hdr1 = None
    close_files(to_close)
    if system == sys_wcs:
        s = StarList(data, columns = ['id', 'ra', 'dec'])
        s_noid = StarList(data_noid, columns=['ra', 'dec'])
    else:
        s = StarList(data, columns = ['id', 'x', 'y'])
        s_noid = StarList(data_noid, columns = ['x', 'y'])
    s.index = s['id']
    s['auto_id'] = False
    if not s_noid.empty:
        id_starts_from = 1 if s.empty else s.id.max() + 1
        ids = range(id_starts_from, id_starts_from + s_noid.stars_number())
        s_noid['id'] = ids
        s_noid.index = ids
        s_noid['auto_id'] = True
        if s.empty:
            s = s_noid
        else:
            s = s.append(s_noid)

    s.DAO_hdr = hdr
    s.DAO_type = DAO.RADEC_FILE if system == sys_wcs else DAO.XY_FILE
    return s


def write_ds9_regions(starlist, filename,
                      color='green', width=1, size=None, font=None, label='{id:.0f}',
                      exclude=None, indexes=None, colors=None, sizes=None, labels=None,
                      color_column=None, size_column=None,
                      comment=None, add_global=None, WCS=False):
    """
    Writes ds9 region file.
    Some regions can be visually distinguish by providing additional indexes to select those regions
    with specific attributes
    :param StarList starlist: StarList object to dump
    :param str filename:      output filename or stream open for writing
    :param str color:         default color
    :param int width:         default line width
    :param int size:          default radius (default 8px or 2")
    :param str font:          ds9 font specification e.g. "times 12 bold italic"
    :param str label:         format expression for label, use col names
    :param pd.Index exclude:  index of disabled regions, if None all are enabled
    :param [pd.Index] indexes: additional indexes to include specific color and size attributes
    :param [str] colors:      specific colors for indexes
    :param [int] sizes:       specific sizes for indexes
    :param [str] labels:      specific labels for indexes
    :param str color_column:  column of starlist with color values
    :param str size_column:   column of starlist with size values
    :param str add_global:    content of additional 'global' if not None
    :param str comment:       content of additional comment line if not None
    :param bool or str WCS:   If true, columns `ra` and `dec` will be used and coord system set to ICRS
                              If nonepmpty string, string will be used as system description
                              If None, False or '', columns 'x','y' will be used and system set to IMAGE
    Example:
    write_ds9_regions(sl, 'i.reg', color='blue',
                        indexes=[saturated, psf],
                        colours=['yellow', 'red'],
                        sizes=[12, None],
                        labels=[None, 'PDF:{id}'],
                        exclude=faint)
    Generates regions file i.reg of blue circles, radius 8,
    objects present in index saturated will have larger yellow circles
    objects present in index psf will be red and labeled with prefix PSF:
    objects present in index faint will be disabled by '-' sign and not displayed by ds9, but can be parsed back
    """
    if WCS:
        xcol = 'ra'
        ycol = 'dec'
        starlist = as_starlist(starlist)
    else:
        xcol = 'x'
        ycol = 'y'
        starlist = as_starlist(starlist, updateskycoord=False)
    try: (starlist[xcol], starlist[ycol])
    except KeyError as e:
        raise KeyError('No coordinate columns ({},{}) in starlist. Check WCS parameter also'.format(xcol, ycol))
    f, to_close = get_stream(filename, 'w')
    f.write('# Region file format: DS9 version 4.0\n')
    if starlist.DAO_hdr is not None:
        write_dao_header(starlist.DAO_hdr, f, '#')
    if comment is not None:
        f.write('#{}\n'.format(comment))
    if color is not None:
        f.write('global color={}\n'.format(color))
    if width is not None:
        f.write('global width={}\n'.format(width))
    if font is not None:
        f.write('global font={}\n'.format(font))
    if add_global is not None:
        f.write('global {}\n'.format(add_global))
    if not WCS:
        f.write('image\n')
    else:
        system = WCS if isinstance(WCS, str) else 'icrs'
        f.write(system+'\n')
    for i, row in starlist.iterrows():
        if exclude is not None and i in exclude:
            f.write('-')
        if size is not None:
            s = size
        else:
            s = '2"' if WCS else 8
        text = label.format(**row)
        c = ''
        if size_column is not None:
            s = row[size_column]
        if color_column is not None:
            c = ' color=' + row[color_column]
        if indexes is not None:
            for n in range(len(indexes)):
                if i in indexes[n]:
                    if sizes and sizes[n] is not None:
                        s = sizes[n]
                    if colors and colors[n] is not None:
                        c = ' color=' + colors[n]
                    if labels and labels[n] is not None:
                        text = labels[n].format(**row)
        f.write('circle({},{},{}) #{} text="{}" id={:d}\n'.format(row[xcol], row[ycol], s, c, text, i))
    close_files(to_close)
