#
# catalog.py -- DSS and star catalog interfaces for the Ginga fits viewer
# 
# Eric Jeschke (eric@naoj.org)
# Raymond Plante -- pyvo interfaces
#
# Copyright (c)  Eric R. Jeschke.  All rights reserved.
# This is open-source software licensed under a BSD license.
# Please see the file LICENSE.txt for details.
#
from __future__ import print_function
import os.path
import tempfile
import re
import urllib
import ginga.util.six as six
if six.PY2:
    from urllib2 import Request, urlopen, URLError, HTTPError
else:
    # python3
    from urllib.request import Request, urlopen
    from urllib.error import URLError, HTTPError
import time

from ginga.misc import Bunch
from ginga.util import wcs


star_attrs = ('name', 'ra', 'dec', 'ra_deg', 'dec_deg', 'mag', 'preference',
              'priority', 'flag', 'b_r', 'dst', 'description')


# Do we have astropy.vo installed?
have_astropy = False
try:
    from astropy.vo.client import conesearch
    from astropy import coordinates, units
    have_astropy = True

except ImportError:
    pass

# How about pyvo?
have_pyvo = False
try:
    import pyvo
    have_pyvo = True

except ImportError:
    pass


class Star(object):
    def __init__(self, **kwdargs):
        starInfo = {}
        starInfo.update(kwdargs)
        for attrname in star_attrs:
            starInfo[attrname] = kwdargs.get(attrname)
        self.starInfo = starInfo

    def __getitem__(self, key):
        return self.starInfo[key]
        
    def __contains__(self, key):
        return key in self.starInfo.keys()
        
    def __setitem__(self, key, value):
        self.starInfo[key] = value
        
    def has_key(self, key):
        return key in self.starInfo
        

class AstroPyCatalogServer(object):

    def __init__(self, logger, full_name, key, url, description):
        self.logger = logger
        self.full_name = full_name
        self.short_name = key
        self.description = description
        self.kind = 'astropy.vo-catalog'
        self.url = url

        # For compatibility with URL catalog servers
        self.params = {}
        count = 0
        for label, key in (('RA', 'ra'), ('DEC', 'dec'), ('Radius', 'r')):
            self.params[key] = Bunch.Bunch(name=key, convert=str,
                                           label=label, order=count)
            count += 1

    def getParams(self):
        return self.params

    def toStar(self, data, ext, magfield):
        try:
            mag = float(data[magfield])
        except:
            mag = 0.0
            
        # Make sure we have at least these Ginga standard fields defined
        d = { 'name':         data[ext['id']],
              'ra_deg':       float(data[ext['ra']]),
              'dec_deg':      float(data[ext['dec']]),
              'mag':          mag,
              'preference':   0.0,
              'priority':     0,
              'description':  'fake magnitude' }
        data.update(d)
        data['ra'] = wcs.raDegToString(data['ra_deg'])
        data['dec'] = wcs.decDegToString(data['dec_deg'])
        return Star(**data)

    def search(self, **params):
        """For compatibility with generic star catalog search.
        """

        self.logger.debug("search params=%s" % (str(params)))
        ra, dec = params['ra'], params['dec']
        if not (':' in ra):
            # Assume RA and DEC are in degrees
            ra_deg = float(ra)
            dec_deg = float(dec)
        else:
            # Assume RA and DEC are in standard string notation
            ra_deg = wcs.hmsStrToDeg(ra)
            dec_deg = wcs.dmsStrToDeg(dec)

        # Convert to degrees for search radius
        radius_deg = float(params['r']) / 60.0
        #radius_deg = float(params['r'])

        c = coordinates.ICRSCoordinates(ra_deg, dec_deg,
                                        unit=(units.degree, units.degree))
        self.logger.info("Querying catalog: %s" % (self.full_name))
        time_start = time.time()
        results = conesearch.conesearch(c, radius_deg * units.degree,
                                        catalog_db=self.full_name)
        time_elapsed = time.time() - time_start

        numsources = results.array.size
        self.logger.info("Found %d sources in %.2f sec" % (
            numsources, time_elapsed))

        # Scan the returned fields to find ones we need to extract
        # particulars from (ra, dec, id, magnitude)
        mags = []
        ext = {}
        fields = results.array.dtype.names
        for name in fields:
            ucd = results.get_field_by_id(name).ucd
            ucd = str(ucd).lower()
            if ucd == 'id_main':
                ext['id'] = name
            elif ucd == 'pos_eq_ra_main':
                ext['ra'] = name
            elif ucd == 'pos_eq_dec_main':
                ext['dec'] = name
            if ('phot_' in ucd) or ('phot.' in ucd):
                mags.append(name)
        self.logger.debug("possible magnitude fields: %s" % str(mags))
        if len(mags) > 0:
            magfield = mags[0]
        else:
            magfield = None

        # prepare the result list
        starlist = []
        arr = results.array
        for i in range(numsources):
            source = dict(zip(fields, arr[i]))
            starlist.append(self.toStar(source, ext, magfield))

        # metadata about the list
        columns = [('Name', 'name'),
                   ('RA', 'ra'),
                   ('DEC', 'dec'),
                   ('Mag', 'mag'),
                   ('Preference', 'preference'),
                   ('Priority', 'priority'),
                   ('Description', 'description'),
                   ]
        # Append extra columns returned by search to table header 
        # TODO: what if not all sources have same record structure?
        # is this possible with VO?
        cols = list(fields)
        cols.remove(ext['ra'])
        cols.remove(ext['dec'])
        cols.remove(ext['id'])
        columns.extend(zip(cols, cols))

        # which column is the likely one to color source circles
        colorCode = 'Mag'
        
        info = Bunch.Bunch(columns=columns, color=colorCode)
        return starlist, info


class PyVOCatalogServer(object):

    def __init__(self, logger, full_name, key, url, description):
        self.logger = logger
        self.full_name = full_name
        self.short_name = key
        self.description = description
        self.kind = 'pyvo-catalog'
        self.url = url

        # For compatibility with URL catalog servers
        self.params = {}
        count = 0
        for label, key in (('RA', 'ra'), ('DEC', 'dec'), ('Radius', 'r')):
            self.params[key] = Bunch.Bunch(name=key, convert=str,
                                           label=label, order=count)
            count += 1

    def getParams(self):
        return self.params

    def toStar(self, data, ext, magfield):
        try:
            mag = float(data[magfield])
        except:
            mag = 0.0
            
        # Make sure we have at least these Ginga standard fields defined
        d = { 'name':         data[ext['id']],
              'ra_deg':       float(data[ext['ra']]),
              'dec_deg':      float(data[ext['dec']]),
              'mag':          mag,
              'preference':   0.0,
              'priority':     0,
              'description':  'fake magnitude' }
        data.update(d)
        data['ra'] = wcs.raDegToString(data['ra_deg'])
        data['dec'] = wcs.decDegToString(data['dec_deg'])
        return Star(**data)

    def search(self, **params):
        """For compatibility with generic star catalog search.
        """

        self.logger.debug("search params=%s" % (str(params)))
        ra, dec = params['ra'], params['dec']
        if not (':' in ra):
            # Assume RA and DEC are in degrees
            ra_deg = float(ra)
            dec_deg = float(dec)
        else:
            # Assume RA and DEC are in standard string notation
            ra_deg = wcs.hmsStrToDeg(ra)
            dec_deg = wcs.dmsStrToDeg(dec)

        # Convert to degrees for search radius
        radius_deg = float(params['r']) / 60.0
        #radius_deg = float(params['r'])
        
        # initialize our query object with the service's base URL
        query = pyvo.scs.SCSQuery(self.url)
        query.ra = ra_deg
        query.dec = dec_deg
        query.radius = radius_deg
        self.logger.info("Will query: %s" % query.getqueryurl(True))

        time_start = time.time()
        results = query.execute()
        time_elapsed = time.time() - time_start

        numsources = len(results)
        self.logger.info("Found %d sources in %.2f sec" % (
            numsources, time_elapsed))

        # Scan the returned fields to find ones we need to extract
        # particulars from (ra, dec, id, magnitude)
        mags = []
        ext = {}
        fields = results.fielddesc()
        for field in fields:
            ucd = str(field.ucd).lower()
            if ucd == 'id_main':
                ext['id'] = field.name
            elif ucd == 'pos_eq_ra_main':
                ext['ra'] = field.name
            elif ucd == 'pos_eq_dec_main':
                ext['dec'] = field.name
            if ('phot_' in ucd) or ('phot.' in ucd):
                mags.append(field.name)
        self.logger.debug("possible magnitude fields: %s" % str(mags))