# db.py -- Line database classes and instances
# Db is the generic class from which Jpl, Cdms, Slap and Local inherit.
# Mlinedb is a container for multiple Db instances. Most of the high-level code
# is meant to work on an mlinedb instance.
# Local is obviously the implementation for a local on disk database (.db). It uses
# sqlite as a backend (see cache.py), but other backends could be plugged in. For
# the sake of consistency with online databases, the code for Local could be moved
# to a local.py file.

import numpy
import os
from linedb import cache,astro
import linedb.progressbar as pb
import sqlite3
from math import log10

NotFoundError = cache.NotFoundError
blankPartfunc = cache.blankPartfunc
origins = cache.origins
timeout = 60
dbVersion = "1.00"

class Db:
   """
   Molecular and atomic line database

   This class defines a molecular and line database that can be
   searched online or offline. Database can be searched online
   through the SLAP protocol, or directly with HTTP POST or GET
   requests. Databases can also be fetched entirely and cached on a
   disk as an a SQLite database for offline searches.

   """

   def __init__(self, url, cache_file, protocol, online = True,
             name = ""):
      """
      Create a database instance

      This function create a database instance which is accessed
      through a given protocol. Supported protocols are "slap",
      "cdms_post" and "jpl_post".

      Arguments:
      url       -- The URL of the database
      protocol   -- The database protocol
      cache_file  -- The name of cache file
      online     -- Search the online database (default True)
      name      -- The name of the database (default "")
      dbout -- where the writes go

      """

      if not protocol in ["slap", "cdms_post", "jpl_post", "catfile", "local"]:
         raise ValueError("Unknown protocol")
      if online:
         if protocol in ["catfile", "local"]:
            raise ValueError("Local db or catfile cannot be online")
         #else:
            # Should try a request to see whether it is available
      else:
         if not protocol in ["catfile", "local"] and not isDbFile(cache_file):
            raise ValueError("Offline cache not available")

      class linedb_data_class():
         """
         Storage class for linedb run time data

         """

      self.url = url
      self.protocol = protocol
      self.cache_file = os.path.expanduser(cache_file)
      self.online = online
      self.data = linedb_data_class()
      self.name = name

   # TODO(mpl): make it so it is all done at init time, and this is not needed anymore. (later).
   def setonline(self, online):
      if online:
         self.online = True
         return
      if not self.protocol in ["slap", "cdms_post", "jpl_post"]:
         raise Exception("Local db or catfile cannot be online")
      if not isDbFile(self.cache_file):
         raise Exception("Offline cache not available, build it first with dbcache")
      self.online = False

   def search(self, fmin, fmax, species, origin, dbsource, energy, einstein):
      """
      Search for lines in the database

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

      """
      raise NotImplementedError("Subclass must implement abstract method")

   def part_function(self, species, origin, dbsource):
      """
      Search for lines in the database

      Arguments:
      species -- the species name (default All)

      """
      raise NotImplementedError("Subclass must implement abstract method part_function")

   def cache(self, fmin = 0, fmax = 2e6, fstep = 1e4, overwrite = False):
      """
      Cache an online database

      Arguments:
      fmin   -- the minimum frequency in MHz
      fmax   -- the maximum frequency in MHz
      fstep  -- the frequency step in MHz between each query

      """

      if self.protocol == "slap":
         raise NotImplementedError("Caching of partition function " \
                "in VO table format is not implemented.")

      c = cache.Cache(self.cache_file)
      c.create(dbVersion, fmin, fmax, overwrite)

      # Add the lines to the database

      species = []
      f = fmin
      if (fstep > fmax - fmin):
         fstep = int(fmax - fmin)

      widgets = ['Downloading lines:     ', pb.Percentage(), ' ', pb.Bar(), ' ',
               pb.ETA(), ' ']
      pbar = pb.ProgressBar(widgets=widgets, maxval= int(fmax - fmin)).start()

      while f < fmax:
         self.online = True # Make sure to search online!
         lines = self.search(fmin = f, fmax = f + fstep + 1e-4)
         for l in lines:
            if not(l.species) in species:
               species.append(l.species)
         c.add_lines(lines, False)
         f += fstep
         pbar.update(f - fmin)

      pbar.finish()


      # Add the partition functions

      widgets = ['Downloading part.func.: ', pb.Percentage(), ' ',
               pb.Bar(), ' ', pb.ETA(), ' ']
      pbar = pb.ProgressBar(widgets=widgets, maxval= len(species)).start()

      # afaik cdms has ~600 species, so it does not seem unreasonable to push all the
      # partfuncs in one call. We could always buffer/group them by 100 if necessary.
      stpo = []
      for i in range(len(species)):
         try:
            temperature, partfunc = self.part_function(species[i], self.name, self.name)
            stpo.append((species[i], temperature, partfunc, self.name, self.name))
#            c.add_partfunc(species[i], temperature, partfunc, self.name, False)
         except NotFoundError:
            continue
         pbar.update(i)
      try:
         c.add_partfuncs(stpo, False)
      except Exception as error:
         raise Exception("While inserting the partfuncs: %s. Database now in inconsistent state." % error)
      pbar.finish()

   def update(self):
      """
      Update the cache file of a database

      """

      pass

   def partition_function(self, T, origin = 'All', dbsource = 'All', *arg, **kw):
      """
      Returns the partition function

      This function returns the partition function of a species at a
      given temperature. The partition function values at different
      temperatures is fetched from the database (online searches) or
      queried in the cache (offline searches). The value of the
      partition function at the requested temperature is obtained by
      a linear fit of the partition function values at the the two
      closest temperature from the desired value.

      Arguments:
      url   -- the URL of partition function file, for votable format
      species -- the name of the species, for other formats and offline searches
      T     -- the temperature

      """

      if (T <= 0.0):
              raise ValueError("Input temperature is negative")
      if self.online and self.protocol == "slap":
         if not "url" in list(kw.keys()):
            raise ValueError("The URL of the partition function must be provided.")
      else:
         if not "species" in list(kw.keys()):
            raise ValueError("The species name must be provided.")

      if self.protocol == "slap":
         urlOrSpecies = kw["url"]
      else:
         urlOrSpecies = kw["species"]
      try:
         temperature, partfunc = self.part_function(urlOrSpecies, origin, dbsource)
      except NotFoundError as error:
         partfunc = [blankPartfunc]
      # TODO(mpl): not sure if we can get a partfunc == [blankPartfunc] and yet no Exception.
      # so I'm just not taking any chances here.
      if partfunc == [blankPartfunc]:
         raise NotFoundError("No partition function for {0} in {1}".format(urlOrSpecies, self.name))

      # Make a linear fit of the two closest values
      temperature = numpy.array(temperature)
      partfunc = numpy.array(partfunc)
      index = temperature.argsort()
      ipos = numpy.searchsorted(numpy.sort(temperature),T)
      if (ipos == 0):
          print("Warning: extrapolated partition function for {0} @ {1} K".format(urlOrSpecies,T))
          i0 = index[0]
          i1 = index[1]
      elif (ipos == len(temperature)):
          print("Warning: extrapolated partition function for {0} @ {1} K".format(urlOrSpecies,T))
          i0 = index[-2]
          i1 = index[-1]
      else:
          i0 = index[ipos-1]
          i1 = index[ipos]
      t0 = temperature[i0]
      t1 = temperature[i1]
      p0 = partfunc[i0]
      p1 = partfunc[i1]
      log10p = (log10(p1)-log10(p0)) / (log10(t1)-log10(t0)) * (log10(T) - log10(t0)) + log10(p0)
      return 10**log10p

   def insertLines(self, dbin, lines, update):
      raise NotImplementedError("Only implemented by Local subclass")

   def removeLines(self, lines):
      raise NotImplementedError("Only implemented by Local subclass")

   def decode_stat_weight(self, stat_weight):
      """
      Decode the statistical weight value

      For some lines, the statistical weight is coded with letter,
      which allow to write number that would not otherwise fit on
      three digit. The letter correspond to the hundreds, i.e. A is
      10, B is 11, etc. Decode the value, and returns an integer.

      Arguments:
      stat_weight -- string (3 chars max) coding the statistical weight
      """

      ascii_code = ord(stat_weight[0])
      if  65 <= ascii_code <= 90:
         hundred_digit = "%s" % (ascii_code - 55)
         stat_weight = hundred_digit + stat_weight[1:]
      return int(stat_weight)


class mlinedb:
   """
   Multiple molecular and atomic line databases

   This class defines a collection of molecular and atomic
   databases. Theses database are searched in sequence using the same
   methods than linedb.

   """

   def __init__(self, dblist):
      """
      Create a multiple database instance

      Arguments:
      dblist -- a list of linedb instances

      """

      self.databases = dblist
      if len(dblist) > 1:
         self.multiple = True
      else:
         self.multiple = False

   def setonline(self, online):

      for db in self.databases:
         db.setonline(online)

   def search(self, fmin, fmax, species, origin, dbsource, energy, einstein):
      """
      Search for lines in the databases

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

      """

      lines = []
#      import cProfile
      for db in self.databases:
         lines += db.search(fmin, fmax, species, origin, dbsource, energy, einstein)
#         cProfile.runctx('db.search(fmin, fmax, species, origin, energy, einstein)', globals(), locals(), '/home/mpl/work/gildas/run/weeds/searchcalls')

      return lines

   def cache(self, fmin = 0, fmax = 2e6, fstep = 1e4, overwrite = False):
      if self.multiple:
         raise ValueError("Cannot cache multiple databases")
      return self.databases[0].cache(fmin, fmax, fstep, overwrite)

   def insertLines(self, dbin, lines, update):
      if self.multiple:
         raise ValueError("Cannot insert on multiple databases")
      return self.databases[0].insertLines(dbin, lines, update)

   def removeLines(self, lines):
      if self.multiple:
         raise ValueError("Cannot remove from multiple databases")
      return self.databases[0].removeLines(lines)

class Local(Db):
   """
   Local sqlite db based on cache.py

   """

   def __init__(self, url, name, overwrite=False):
#      super(Local, self).__init__(url, cache_file="", protocol="local", online=False, name=name)
      Db.__init__(self, url, cache_file="", protocol="local", online=False, name=name)
      self.backend = cache.Cache(url)
      if overwrite:
         self.backend.create(dbVersion, None, None, True)
      else:
         if not isDbFile(url):
            self.backend = None
            raise Exception("%s is not a valid sqlite db" % dbout)

   def search(self, fmin, fmax, species, origin, dbsource, energy, einstein):
      """
      Search lines in the cache

      Arguments:
      fmin   -- the minimum frequency in MHz
      fmax   -- the maximum frequency in MHz
      species -- the species name (a string or list of strings, default no search for species)
      Eu_max -- maximum upper level energy expressed
              in cm-1 (default None)

      """

      lines = self.backend.search(fmin, fmax, species, origin, dbsource, energy, einstein)
      # N.B: can't set it directly within cache.search() because Cache class
      # only knows the full url
      for l in lines:
         l.dbsource = self.name
      return lines

   def cache(self, fmin = 0, fmax = 2e6, fstep = 1e4, overwrite = False):
      raise NotImplementedError("Already a local db")

   def part_function(self, species, origin, dbsource):
      return self.backend.partition_function(species, origin, dbsource)

   # TODO: yet to be tested with a slap db that supports partition_function_link
   # Note that regardless of what options.update is set to, a line will be updated
   # if it's coming from one of the online sources, because we want to keep track
   # of the date change in that case.
   def insertLines(self, dbin, lines, update):
      if len(lines) <= 0:
         return

      c = self.backend

      # Add the lines to the database
      try:
         c.add_lines(lines, update)
      except Exception as error:
         raise Exception("While inserting lines: %s. Database now in inconsistent state." % error)

      # track down all the (species, origin, dbsource) triplets
      sos = []
      for l in lines:
         # TODO(mpl): do better. maybe take care of that directly in part_function.
         # quickhack for slap style databases, because part_function wants a url as arg for them
         try:
            l.partition_function_link
            spec = l.species + "|" + l.partition_function_link
         except:
            spec = l.species
         triplet = [spec, l.origin, l.dbsource, l.prev]
         # we could probably use a set instead of having to check if in list
         if not(triplet) in sos:
            sos.append(triplet)

      # index the input dbs by name, for easy referencing right after
      dbBySource = {}
      for d in dbin.databases:
         dbBySource[d.name] = d

      stpod = []
      for spec, ori, dbsrc, prev in sos:
         try:
            fields = spec.split("|")
            species = fields[0]
            if len(fields) > 1:
               specOrUrl = fields[1]
            else:
               specOrUrl = species
            if not prev:
               prev = dbsrc
            temperature, partfunc = dbBySource[dbsrc].part_function(specOrUrl, ori, prev)
         except NotFoundError as error:
            partfunc = [blankPartfunc]
            temperature = [blankPartfunc]
         except Exception as error:
            raise Exception("While getting the partfuncs: %s. Database now in inconsistent state." % error)

         stpod.append((species, temperature, partfunc, ori, dbsrc))
      try:
         c.add_partfuncs(stpod, update)
      except Exception as error:
         raise Exception("While inserting the partfuncs: %s. Database now in inconsistent state." % error)

   def removeLines(self, lines = None):
      c = self.backend
      c.remove(lines)
      # N.B: this is easier and (probably) more efficient than a trigger that
      # would fire for each delete in the line table. However, it would not
      # be ok if we had any concurrency. We would at least need these two
      # to be atomic together.
      # This is not used anymore. there's a trigger now.
      # c.updatePartfuncOnLineDelete()

# wrapper here because main.py should not have to import cache just for that one
def isDbFile(dbfile):
   return cache.isDbFile(dbfile)

class AstroCatalog:
   """
   Support for line DB exported under Astro catalog format
   """

   def __init__(self, url, name, overwrite=False):
      self.url = url
      self.name = name
      self.backend = astro.Astro(url)
      self.backend.create(dbVersion, None, None, overwrite)

   def search(self, fmin, fmax, species, origin, dbsource, energy, einstein):
      raise NotImplementedError("Can not search in an AstroCatalog")

   def cache(self, fmin = 0, fmax = 2e6, fstep = 1e4, overwrite = False):
      raise NotImplementedError("No sense for an AstroCatalog")

   def part_function(self, species, origin, dbsource):
      raise NotImplementedError("No sense for an AstroCatalog")

   def insertLines(self, dbin, lines, update):
      if len(lines) <= 0:
         return

      c = self.backend

      # Add the lines to the database
      try:
         c.add_lines(lines, update)
      except Exception as error:
         raise Exception("While inserting lines: %s." % error)

   def removeLines(self, lines = None):
      raise NotImplementedError("Can not remove lines from an AstroCatalog")
