# slap.py -- linedb implementation for votables type dbs, using the slap protocol

import os
import urllib.request, urllib.parse, urllib.error
import urllib.request, urllib.error, urllib.parse
import math
from linedb import db
from linedb import line
from linedb import votable
from linedb.consts import *
from datetime import datetime

maxFreqGHz = 2000

class Slap(db.Db):

   def search(self, fmin, fmax, species = 'All', origin = 'All', dbsource = 'All', energy = -1, einstein = -1):
      """
      Search lines in a database using SLAP

      This function make a query in a VO-compliant database using the
      Simple Line Access Protocol (SLAP).

      Arguments:
      fmin   -- the minimum frequency in MHz
      fmax   -- the maximum frequency in MHz
      species -- the species name (default All)
      energy -- maximum upper level energy expressed
              in cm-1 (default None)

      """

      lines = []
      if origin != 'All' and origin.lower() != self.name:  # Case-insensitive
         return
      if dbsource != 'All' and dbsource != self.name:
         return

      # TODO(mpl): check FREQUENCY can be used for all slap style dbs
      # Convert the frequencies into wavelenghts
      #wlmax = speed_of_light / (fmin * 1e6)
      #wlmin = speed_of_light / (fmax * 1e6) # m
      if fmin < 0:
         fmin = 0
      if fmax < 0:
         fmax = maxFreqGHz

      # Construct the URL for the database query
      # FixMe: Not all database support the CHEMICAL_ELEMENT
      # keyword. Filter this by hand?

      #slap_url = self.url + "WAVELENGTH=%.9e/%.9e" % (wlmin, wlmax)
      slap_url = self.url + "FREQUENCY=%.9e/%.9e" % (fmin, fmax)
      if species != "All":
         slap_url = slap_url + "&CHEMICAL_ELEMENT=%s" % urllib.parse.quote_plus(species)

      slap_url = slap_url + "&REQUEST=queryData"

      # Get the results in the VO table format and parse it

      try:
         response = urllib.request.urlopen(slap_url, timeout = db.timeout)
      except Exception as error:
         raise Exception("Could not connect to database: %s" % error)

      try:
         vot = votable.VOTable()
         vot.parse(response)

         # Find-out what fields are in the VO table. Each field is
         # identified with its utype or its name. Note that the utype
         # is defined by the SLAP standard, but the name may vary from
         # one database to the other.

         indexes = {}
         fields = vot.getFieldsAttrs()
         for i in range(len(fields)):
            if 'utype' in fields[i]:
               field_id = fields[i]['utype']
            elif 'name' in fields[i]:
               field_id = fields[i]['name']
            else:
               continue

            # Patch for voparis wrong utypes
            if field_id == "ldm:Line.wavelength":
               field_id = "ssldm:Line.wavelength.value"
            if field_id == "ldm:Line.einsteinA" and \
                  fields[i]['name'] == "log10_A_Einstein_coefficient":
               continue # ignore Einstein A in log units
            if field_id.split(":", 1)[0] == "ldm":
               field_id = "ssldm:" + field_id.split(":", 1)[1]

            indexes[field_id] = i

         # Read each row of the VO table an fill the line object.

         for r in vot.getDataRows():

            l = line.line()

            # Mandatory fields
            if "ssldm:Line.species.name" in indexes:
               l.species = vot.getData(r)[indexes["ssldm:Line.species.name"]]
            elif "ssldm:Line.initialElement.name" in indexes:
               l.species = vot.getData(r)[indexes["ssldm:Line.initialElement.name"]]
            elif "ssldm:Line.title" in indexes:
               l.species = vot.getData(r)[indexes["ssldm:Line.title"]]
            else:
               raise Exception("missing species name")
            if "ssldm:Line.wavelength.value" in indexes:
               wl = float(vot.getData(r)[indexes["ssldm:Line.wavelength.value"]])
               l.frequency = speed_of_light / wl * 1e-6 # MHz
            else:
               raise Exception("missing wavelenght")
            l.origin = self.name
            l.dbsource = self.name
            l.date = datetime.utcnow().isoformat()

            # Optional fields
            if "ssldm:Line.einsteinA" in indexes:
               if vot.getData(r)[indexes["ssldm:Line.einsteinA"]]: # check for an empty field
                  l.einstein_coefficient = float(vot.getData(r)[indexes["ssldm:Line.einsteinA"]])
                  # Some databases returns the Einstein A in log
                  # units, try to guess if that's the case
                  if l.einstein_coefficient < 0:
                     l.einstein_coefficient = 10**l.einstein_coefficient
            #TODO: What about the lower_level.energy? ask the LAOG guys.
            if "ssldm:Line.initialLevel.energy" in indexes:
               if vot.getData(r)[indexes["ssldm:Line.initialLevel.energy"]]:
                  l.upper_level.energy = float(vot.getData(r)[indexes["ssldm:Line.initialLevel.energy"]])
            energyK = l.upper_level.energy * cm_K
            # filter by upper level energy, if required
            if energy > 0 and energyK > energy:
               continue
            # filter by einstein coefficient, if required
            if einstein is not None and einstein > 0 and l.einstein_coefficient < einstein:
               continue

            if "ssldm:Line.initialLevel.name" in indexes:
               if vot.getData(r)[indexes["ssldm:Line.initialLevel.name"]]:
                  l.upper_level.quantum_numbers = vot.getData(r)[indexes["ssldm:Line.initialLevel.name"]]
            if "ssldm:Line.finalLevel.name" in indexes:
               if vot.getData(r)[indexes["ssldm:Line.finalLevel.name"]]:
                  l.lower_level.quantum_numbers = vot.getData(r)[indexes["ssldm:Line.finalLevel.name"]]
            elif "quantum numbers" in indexes:
               # Splatalogue does not separate the quantum numbers
               # from the upper and lower levels
               l.upper_level.quantum_numbers = vot.getData(r)[indexes["quantum numbers"]]
            if "frequency recommended" in indexes:
               # In Splatalogue this field is set to 1 is the
               # frequency is recommended by the NRAO (e.g. if one of
               # catalog -- JPL or CDMS is known to provide more
               # accurate frequency than the other). I think this
               # also selects known/probable interstellar species,
               # but this need to be checked.
               if vot.getData(r)[indexes["frequency recommended"]]:
                  l.recommended = int(vot.getData(r)[indexes["frequency recommended"]])
               else:
                  l.recommended = None
               if l.recommended != 1:
                  continue # ignore not-recommended values

            # The following fields are not defined by the SLAP
            # standard, but may be returned by some databases (and are
            # needed for LTE modeling)
            if "partition_function_link" in indexes:
               l.partition_function_link = vot.getData(r)[indexes["partition_function_link"]]
            else:
               l.partition_function_link = db.blankPartfunc
            if "initial_statistical_weight" in indexes:
               l.upper_level.statistical_weight = float(vot.getData(r)[indexes["initial_statistical_weight"]])
            if "final_statistical_weight_index" in indexes:
               l.lower_level.statistical_weight = float(vot.getData(r)[indexes["final_statistical_weight_index"]])

            lines.append(l)

      except Exception as error:
         raise Exception("Can't parse the response from the database: %s" % error)

      return lines

   def part_function(self, url, origin, dbsource):
      """
      Returns the partition function at different temperatures

      This function fetches a VO table containing the partition function
      at different temeperatures from the given URL, and then parse it.

      Arguments:
      url -- the URL of the VO table

      """

      if origin.lower() != self.name:  # Case-insensitive
         raise ValueError("Got %s, but want slap as origin for partfunc in slap" % origin)
      if dbsource != self.name:
         raise ValueError("Got %s, but want slap as dbsource for partfunc in slap" % dbsource)

      if url == db.blankPartfunc:
         raise db.NotFoundError("No partition_function_link")

      vot = votable.VOTable()
      vot.parse (urllib.request.urlopen(url))

      temperature_index = vot.getColumnIdx('temperature')
      value_index = vot.getColumnIdx('value')

      temperature = []
      partfunc = []

      for r in vot.getDataRows():
         temperature.append(float(vot.getData(r)[temperature_index]))
         partfunc.append(float(vot.getData(r)[value_index]))

      if len(partfunc) == 0:
         raise NotFoundError("No partition function found for this species.")

      return temperature, partfunc
