# jpl.py -- linedb implementation for jpl type dbs
# Jpl class is obviously for online jpl. Catfile class is for local .cat files,
# following the jpl format, and using partfunc.cat or partfunc.json for the
# partition functions.
# JplReader class holds code common to both Jpl and Catfile. In retrospect this
# design might not be the most adequate and it may be cleaner to completely decouple
# Catfile from the jpl format.

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

if 'GAG_LINEDB_JPL_PARTFUNCURL' in os.environ:
  catdir_jpl_url = os.environ['GAG_LINEDB_JPL_PARTFUNCURL']
else:
  catdir_jpl_url = "https://spec.jpl.nasa.gov/ftp/pub/catalog/catdir.cat"

maxFreqGHz = 2000
qpart = "qpart"
ptemp = [300., 225., 150., 75., 37.5, 18.75, 9.375]

class JplReader(db.Db):
   """ This class is an attempt to factorize code which is used
   both when working with the remote jpl db, and with data from
   a local .cat file (which uses the jpl format).
   """

   def parse_line(self, l, fmin, fmax, species, energy, einstein, tg):
      """ Parse the line string l and derive all the line properties.
      See description here: https://spec.jpl.nasa.gov/ftp/pub/catalog/README

      Filters (if provided) are also applied, namely by:
        - frequency range
        - species name(s): list of desired species (empty list means select all)
        - energy level
        - einstein coefficient
      Return nothing if the line is filtered out.
      """
      species_tag = self.data.species_tag
      tag_species = self.data.tag_species
      species_Q = self.data.species_Q
      freq = float(l[0:13]) # MHz
      # TODO(mpl): this test is unnecessary when parsing a response from
      # __post, maybe we can avoid it.
      # Yes, but this code is also used for searching in offline ASCII
      # databases. Beware. (SB, 25-mar-2014)
      if (fmin > 0 and freq < fmin) or (fmax > 0 and freq > fmax):
         return
      wavelength = speed_of_light / (freq * 1e6) * 1e2 # cm
      errfreq = float(l[13:21])  # MHz
      log_intensity = float(l[21:29])
      lower_level_energy = float(l[31:41]) * cm_K # K
      upper_level_statistical_weight = self.decode_stat_weight(l[41:44])
      upper_level_quantum_numbers = l[55:67].rstrip().lstrip()
      lower_level_quantum_numbers = l[67:].rstrip().lstrip()

      # Get the species name and partition function
      spec = tag_species[tg]
      if len(species)>0:
        # Filter out this species as required
        if spec not in species:
          return
      try:
         Q = species_Q[tg][300.]
      except KeyError as error:
         raise Exception("partition function at 300.K is missing".format(error))

      # Compute the Einstein coefficient from the line intensity
      # at 300 K. For the formulae used see Eq. 5 on:
      # http://www.astro.uni-koeln.de/cdms/catalog
      upper_level_energy = lower_level_energy + cm_K / wavelength # K
      # filter by upper level energy, if required
      if energy > 0 and upper_level_energy > energy:
         return

      einstein_coefficient = 2.7964e-16 * 10**log_intensity * freq**2 \
         * Q / upper_level_statistical_weight \
         / (math.exp(-lower_level_energy / 300.) - math.exp(-upper_level_energy / 300.))
      # filter by einstein coefficient, if required
      if einstein is not None and einstein > 0 and einstein_coefficient < einstein:
         return

      sl = line.line()
      sl.species = spec
      sl.frequency = freq
      sl.err_frequency = errfreq
      sl.einstein_coefficient = einstein_coefficient
      sl.upper_level.energy = upper_level_energy
      sl.upper_level.statistical_weight = upper_level_statistical_weight
      sl.upper_level.quantum_numbers = upper_level_quantum_numbers
      sl.lower_level.energy = lower_level_energy
      sl.lower_level.statistical_weight = upper_level_statistical_weight
      sl.lower_level.quantum_numbers = lower_level_quantum_numbers
      sl.origin = self.name
      sl.dbsource = self.name
      sl.date = datetime.utcnow().isoformat()
      return sl

class Jpl(JplReader):

   def __get_context(self):
      # The method urllib2.urlopen is able to deal with secured
      # connections, but unfortunately the protocol it uses by default
      # under Mac (unclear which one) is not high enough for JPL which
      # enforced only the strongest ones (TLS v1.1 and v1.2, it
      # rejects older ones).
      try:
        # PROTOCOL_TLS selects the highest protocol version that both the
        # client and server support. Available since Python 3.6
        import ssl
        return ssl.SSLContext(ssl.PROTOCOL_TLS)
      except:
        pass
      try:
        # Define a SSL context with highest TLS
        import ssl
        return ssl.SSLContext(ssl.PROTOCOL_TLSv1_2)
      except:
        # Context failed... let it undefined...
        return None

   def __read_catdir(self, url):
      # In the JPL database, species are identified by a tag, so we must
      # first first find what is the tag of each species. We can get
      # these from the catdir.cat file. This file also contains the
      # partition function values at 300 K which we need to compute the
      # Einstein coefficient of each line.

      # No need to get it again if our internal dicts are already setup
      if (hasattr(self.data, 'species_tag') and hasattr(self.data, 'tag_species') and hasattr(self.data, 'species_Q')):
         return

      try:
         context = self.__get_context()
         if context is None:
           # Older version of urlopen do not understand "context=" argument
           # even if it is None. Skip it.
           response =  urllib.request.urlopen(url)
         else:
           response =  urllib.request.urlopen(url,context=context)
      except:
         print("failed to open url: %s" % url)
         raise

      species_tag, tag_species = {}, {}
      species_Q = {}

      for l in response.readlines():
         l = l.decode('utf-8')
         try:
            qs = []
            tag = int(l[0:6])
            s = l[7:20].rstrip()
            qs.append(10**float(l[26:33])) # Partition function at 300 K
            # caching the other usual partfunc values right here and right now
            # for later use by part_function
            qs.append(10**float(l[33:40]))
            qs.append(10**float(l[40:47]))
            qs.append(10**float(l[47:54]))
            qs.append(10**float(l[54:61]))
            qs.append(10**float(l[61:68]))
            qs.append(10**float(l[68:75]))
         except (IndexError, TypeError) as details:
            raise Exception("Error while reading catalog file: %s" % details)
         species_tag[s] = tag
         tag_species[tag] = s
         species_Q[tag] = dict([(t, q) for t, q in zip(ptemp, qs)])
      response.close()

      # Store the results to speed-up future searches
      self.data.species_tag = species_tag
      self.data.tag_species = tag_species
      self.data.species_Q = species_Q

   def __parse_response(self, response, fmin, fmax, species, energy, einstein):
      tg = None
      lines = []
      for l in response.readlines():
         # See comment in cdms.py about string decoding to utf-8
         l = l.decode('utf-8')
         try:
            if (tg is None):
               # No tag found yet
               if 'No catform lines were found' in l:
                  # As of 17-oct-2014, JPL returns the following message:
                  # "No catform lines were found for your search criteria..."
                  return []
               if 'Zero lines were found' in l:
                  # As of 29-jan-2019, JPL returns the following message:
                  # "Zero lines were found for your search criteria."
                  return []

            # Get the species name and tag, and iterate to next line.
            # Line can be e.g. "              47003  DCOOH"
            if len(l) < 36:
              tg = int(l.split()[0])
              continue

            # If the search is incomplete, split the search frequency
            # range in several divisions.
            # FixMe: maybe we should make sure that the lines at the
            # edge of the two frequency range don't appear twice in
            # the list because of rounding.
            # FixMe: this could be implemented at an higher level --
            # this might be useful for other databases too.
            if l[0:9] == "THIS form":
               ndiv = 4
               frange = (fmax-fmin)/ndiv
               lines = []
               f1 = fmin
               for idiv in range(ndiv):
                  f2 = f1+frange
                  lines += self.post(f1, f1+frange,
                                     species, energy, einstein)
                  f1 = f2
               return lines

            if tg is None:
               # We reached this point without defining the tag: file is malformed
               raise Exception("Malformed catalog file: expected a first line with tag and species name before the associated line list")

            # Select all species [] because we already POST'ed the desired
            # selection for this online database
            sl = self.parse_line(l, fmin, fmax, [], energy, einstein, tg)
            if not sl:
               continue
            lines.append(sl)
         except KeyError as error:
            raise Exception("KeyError caught: {0}, which means your local catdir might be outdated".format(error))
         except Exception as error:
            raise Exception("Can't parse the response from the database: %s" % error)
      return lines

   def post(self, fmin, fmax, species, energy, einstein):
      # set up the tag, species, and Q dicts
      try:
         self.__read_catdir(catdir_jpl_url)
      except Exception as error:
         raise Exception("Couldn't read catdir: {0}".format(error))


      # Make a HTTP/POST query on the server. Note that for the JPL
      # database, the keyword order does matter, so we must use a tuple
      # for the form values.

      if fmin > 0:
         fmin_s = "%.9f" % (fmin * 1e-3) # MHz -> GHz
      else:
         fmin_s = 0
      if fmax > 0:
         fmax_s = "%.9f" % (fmax * 1e-3) # MHz -> GHz
      else:
         fmax_s = maxFreqGHz

      if species != "All":
         if species in list(self.data.species_tag.keys()):
            tag = self.data.species_tag[species]
         else:
            return [] # species is not in the database
      else:
         tag = "All"

      form_values = (('MinNu', fmin_s),   ('MaxNu', fmax_s),
                     ('MaxLines',100000), ('UnitNu', "GHz"),
                     ('Mol', tag),        ('StrLim', -500))

      try:
         if 'GAG_LINEDB_JPL_QUERYURL' in os.environ:
            search_url = os.environ['GAG_LINEDB_JPL_QUERYURL']
         else:
            search_url = self.url + "/cgi-bin/catform"
         data = urllib.parse.urlencode(form_values)
         bdata = data.encode('utf-8')
         req = urllib.request.Request(search_url, bdata)
         context = self.__get_context()
         if context is None:
           # Older versions of urlopen do not understand "context=" argument
           # even if it is None. Skip it.
           response = urllib.request.urlopen(req, timeout=db.timeout)
         else:
           response = urllib.request.urlopen(req, timeout=db.timeout, context=context)
      except Exception as error:
         raise Exception("Could not connect to database (%s): %s" % (search_url,error))

      # Parse the results
      lines = self.__parse_response(response, fmin, fmax, species, energy, einstein)
      response.close()
      return lines

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

      Arguments:
      fmin   -- the minimum frequency in MHz
      fmax   -- the maximum frequency in MHz
      species -- the species name (single string). String 'All' is an alias for
                 no selection. Tolerate a list as input with 0 (no selection
                 i.e. same as 'All') or 1 string element.
      energy -- maximum upper level energy expressed
              in cm-1 (default None)
      """

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

      if dbsource != 'All' and dbsource != self.name:
         return lines

      if (type(species) == str):
        lspecies = species
      elif (type(species) == list):
        if (len(species) == 0):
          lspecies = 'All'
        elif (len(species) == 1):
          lspecies = species[0]
        else:
          raise Exception("Selection with several species is not available for online JPL")
          # Actually we could do it by searching all species and then using the species
          # filter in the parse_line() method.
      else:
        raise Exception("Unexpected kind of argument: "+repr(species))

      if self.protocol != "jpl_post":
         raise ValueError("Unsuitable protocol: {0}".format(self.protocol))

      if self.online:
         lines = self.post(fmin, fmax, lspecies, energy, einstein)
      else:
         raise Exception("Offline in jpl instance")

      return lines

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

      This function fetches a text file (in the format used in the JPL
      database) containing the partition function of the given species
      for different temperatures and then parses it. The file content is
      kept in memory to avoid re-fetching if the function is called
      several times.

      Arguments:
      species -- the species name

      """

      if self.protocol != "jpl_post":
         raise ValueError("Unsuitable protocol: {0}".format(self.protocol))

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

      if dbsource != self.name:
         raise ValueError("Got %s, but want jpl as dbsource for partfunc in jpl" % dbsource)

      try:
         self.__read_catdir(catdir_jpl_url)
      except Exception as error:
         raise Exception("Couldn't read catdir: {0}".format(error))
      try:
         tempToQ = self.data.species_Q[self.data.species_tag[species]]
      # do we really want a db.NotFoundError ?
      except KeyError:
         raise db.NotFoundError("No partition function found for %s." % species)
      return ptemp, [tempToQ[t] for t in ptemp]


default = Jpl(url = "https://spec.jpl.nasa.gov",
          cache_file = "~/.gag/scratch/jpl.db", protocol = "jpl_post",
          online = True, name = "jpl")

class Catfile(JplReader):

   def __local_partfunc_url(self):
      # assume for now that the partfunc file is located in the same dir as
      # the input file, and is named either partfunc.json or partfunc.cat
      dir, dontcare = os.path.split(self.url)
      pfjson = os.path.join(dir, "partfunc.json")
      if os.path.isfile(pfjson):
         return pfjson
      else:
         return os.path.join(dir, "partfunc.cat")

   def __read_partfunc(self, filepath):
      # No need to get it again if our internal dicts are already setup
      if (hasattr(self.data, 'species_tag') and hasattr(self.data, 'tag_species') and hasattr(self.data, 'species_Q')):
         return
      dontcare, ext = os.path.splitext(filepath)
      if ext == ".json":
         f = self.__read_partfuncjson
      else:
         f = self.__read_partfuncat
      f(filepath)

   def __read_partfuncjson(self, filepath):
      try:
         response = open(filepath, "r")
         data = json.load(response)
      except Exception as error:
         raise Exception("loading {0}: {1}".format(filepath, error))

      species_tag, tag_species = {}, {}
      species_Q = {}
      try:
         for d in data:
            try:
               st = d['species']
            except:
               raise Exception("'species' key missing")
            if len(st) != 2:
                raise Exception("Wrong size for 'species' list")
            spe = st[0]
            tag = int(st[1])
            try:
               temp = d['temperatures']
            except:
               raise Exception("'temperatures' key missing for {0}".format(s))
            if len(temp) < 1:
                raise Exception("empty temperatures list for {0}".format(s))
            try:
               Q = d[qpart]
            except:
               raise Exception("qpart key missing for {0}".format(s))
            if len(Q) != len(temp):
                raise Exception("qpart list size different from temperature list size for {0}".format(s))
            tempToQ = dict((t, 10**q) for t, q in zip(temp, Q))
            species_tag[spe] = tag
            tag_species[tag] = spe
            species_Q[tag] = tempToQ
      except Exception as error:
         raise Exception("Reading json partfunc index: {0}".format(error))
      response.close()
      # Store the results to speed-up future searches
      self.data.species_tag = species_tag
      self.data.tag_species = tag_species
      self.data.species_Q = species_Q

   def __read_partfuncat(self, filepath):
      try:
         response = open(filepath, "r")
      except Exception as error:
         raise Exception("opening {0}: {1}".format(filepath, error))

      species_tag, tag_species = {}, {}
      species_Q = {}

      iline = 0
      while True:
         try:
            # Species
            iline += 1
            l = response.readline()
            if len(l) == 0:
               break
            if len(l.strip()) == 0:
               continue
            if l[0] == "#":
               continue
            fields = shlex.split(l)  # Takes care of blanks in double-quoted strings
            if fields[0].lower() != "species":
               message = "\nExpected tag 'species' at line {0}, got:\n{1}".format(iline,l)
               raise Exception(message)
            elif len(fields) != 3:
               message = "\n'species' tag expects 2 arguments at line {0}, got:\n{1}".format(iline,l)
               raise Exception(message)
            spe = fields[1]
            tag = int(fields[2])

            # Temperatures
            iline += 1
            l = response.readline()
            fields = l.split()
            if fields[0].lower() != "temperatures":
               message = "\nExpected tag 'temperatures' at line {0}, got:\n{1}".format(iline,l)
               raise Exception(message)
            elif len(fields) < 2:
               message = "\nThere must be at least 1 temperature at line {0}, got:\n{1}".format(iline,l)
               raise Exception(message)
            nf = len(fields)
            temps = [float(t) for t in fields[1:]]

            # Partfunc
            iline += 1
            l = response.readline()
            fields = l.split()
            if fields[0].lower() != qpart:
               message = "\nExpected tag '{0}' at line {1}, got:\n{2}".format(qpart,iline,l)
               raise Exception(message)
            elif len(fields) != nf:
               message = "\nNumber of elements differ in 'temperatures' and '{0}' tags at line {1}, got:\n{2}".format(qpart,iline,l)
               raise Exception(message)
            tempToQ = dict((t, 10**float(q)) for t, q in zip(temps, fields[1:]))

         except Exception as error:
            raise Exception("Reading partfunc index: {0}".format(error))
         species_tag[spe] = tag
         tag_species[tag] = spe
         species_Q[tag] = tempToQ
      response.close()

      # Store the results to speed-up future searches
      self.data.species_tag = species_tag
      self.data.tag_species = tag_species
      self.data.species_Q = species_Q

   def __parse_response(self, response, fmin, fmax, species, energy, einstein):
      tg = None
      lines = []
      for l in response.readlines():
         try:
            # Skip lines containing the species name and tag
            if len(l) < 36:
               args = shlex.split(l)  # Takes care of blanks in double-quoted strings
               tg = int(args[0])
               # Sanity check
               if (args[1] != self.data.tag_species[tg]):
                 raise Exception('Species name mismatch for tag %d: "%s" vs "%s"' % (tg,args[1],self.data.tag_species[tg]))
               continue
            if tg is None:
               # We reached this point without defining the tag: file is malformed
               raise Exception("Malformed catalog file: expected a first line with tag and species name before the associated line list")

            sl = self.parse_line(l, fmin, fmax, species, energy, einstein, tg)
            if not sl:
               continue
            lines.append(sl)
         except KeyError as error:
            raise Exception("KeyError caught: {0}, which means your local catdir might be outdated".format(error))
         except Exception as error:
            raise Exception("Can't parse the response from the database: %s" % error)
      return lines

   def __catfile(self, fmin, fmax, species=[], energy=-1, einstein=-1):
      """
      Search for lines in a local catalog

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

      """
      # set up the tag, species, and Q dicts
      try:
         self.__read_partfunc(self.__local_partfunc_url())
      except Exception as error:
         raise Exception("catfile search: {0}".format(error))

      response = open(self.url, "r")

      if (type(species) == str):
        if (species == 'All'):
          lspecies = []  # Empty list (no selection by species)
        else:
          lspecies = [species]  # Store single string to a list
      elif (type(species) == list):
        lspecies = species
      else:
        raise Exception("Unexpected kind of argument: "+repr(species))

      # Parse the results
      lines = self.__parse_response(response, fmin, fmax, lspecies, energy, einstein)
      response.close()
      return lines

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

      Arguments:
      fmin   -- the minimum frequency in MHz
      fmax   -- the maximum frequency in MHz
      species -- the specie 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
      if self.protocol != "catfile":
         raise ValueError("Unsuitable protocol: {0}".format(self.protocol))
      if self.online:
         raise Exception("Catfile cannot be online")

      lines = self.__catfile(fmin, fmax, species, energy)
      return lines

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

      This function fetches a text file (in the format used in the JPL
      database) containing the partition function of the given species
      for different temperatures and then parses it. The file content is
      kept in memory to avoid re-fetching if the function is called
      several times.

      Arguments:
      species -- the species name

      """
      if self.protocol != "catfile":
         raise ValueError("Unsuitable protocol: {0}".format(self.protocol))
      if origin.lower() != self.name:  # Case-insensitive
         raise ValueError("Got {0}, but want {1} as origin for partfunc in jpl".format(origin, self.name))
      if dbsource != self.name:
         raise ValueError("Got {0}, but want {1} as dbsource for partfunc in jpl".format(dbsource, self.name))
      return self.part_functionT(species, [])

   def part_functionT(self, species, temps):
      filepath = self.__local_partfunc_url()
      dontcare, ext = os.path.splitext(filepath)
      if ext == ".json":
         f = self.__part_functionTjson
      else:
         f = self.__part_functionTcat
      return f(species, temps)

   def __part_functionTjson(self, species, temps):
      filepath = self.__local_partfunc_url()
      self.__read_partfuncjson(filepath)
      try:
         tempToQ = self.data.species_Q[self.data.species_tag[species]]
      # do we really want a db.NotFoundError ?
      except KeyError:
         raise db.NotFoundError("No partition function found for %s." % species)
      if temps == []:
         T = sorted(tempToQ.keys())
         Q = [tempToQ[t] for t in T]
         return T, Q
      try:
         Q = [tempToQ[t] for t in temps]
      except KeyError as error:
         raise db.NotFoundError("partfunc not found for species {1}: {0}".format(error, species))
      return temps, Q

   def __part_functionTcat(self, species, temps):
      filepath = self.__local_partfunc_url()
      self.__read_partfuncat(filepath)
      try:
         tempToQ = self.data.species_Q[self.data.species_tag[species]]
      # do we really want a db.NotFoundError ?
      except KeyError:
         raise db.NotFoundError("No partition function found for %s." % species)
      if temps == []:
         T = sorted(tempToQ.keys())
         Q = [tempToQ[t] for t in T]
         return T, Q
      try:
         Q = [tempToQ[t] for t in temps]
      except KeyError as error:
         raise db.NotFoundError("partfunc not found for species {1}: {0}".format(error, species))
      return temps, Q

