
import pysic
from operator import attrgetter
#from functools import cmp_to_key
import linedb

def select(command, fmin, fmax, species, origin, dbsource, energy, einstein,
   sortby=None,s_epsilon=-1):
   """
   Return a list of lines according to the selection criteria. Search is
   made in the current input database (selected with USE IN).

   """

   # Search the database for the lines of the given species in the
   # frequency range.

   if (linedb.dbin is None):
     pysic.message(pysic.seve.e,command,"No database selected")
     pysic.sicerror()
     return []

   try:
      output = linedb.dbin.search(fmin, fmax, species, origin, dbsource, energy, einstein)
   except Exception as error:
      raise

   if len(output) == 0:
      severity = pysic.seve.w
      message = "Nothing found"
   else:
      severity = pysic.seve.i
      min = fmin
      max = fmax
      if fmin < 0:
         min = 0
      if fmax < 0:
         max = 'infinity'
      message = "{0} lines found in the frequency range {1} to {2} MHz".format(len(output), min, max)
   #
   if (origin != 'All'):
     message += " (origin "+origin+")"
   pysic.message(severity,command,message)
   #
   # Sort if requested
   sortLines(output,sortby,s_epsilon)
   #
   return output


# cmp_to_key is provided here instead of being imported from functools,
# because it is only shipped with python >= 2.7
def cmp_to_key(mycmp):
   'Convert a cmp= function into a key= function'
   class K(object):
      def __init__(self, obj, *args):
         self.obj = obj
      def __lt__(self, other):
         return mycmp(self.obj, other.obj) < 0
      def __gt__(self, other):
         return mycmp(self.obj, other.obj) > 0
      def __eq__(self, other):
         return mycmp(self.obj, other.obj) == 0
      def __le__(self, other):
         return mycmp(self.obj, other.obj) <= 0
      def __ge__(self, other):
         return mycmp(self.obj, other.obj) >= 0
      def __ne__(self, other):
         return mycmp(self.obj, other.obj) != 0
   return K


def sortLines(lines,criterion,epsilon):
   """
   Sorts in place the 'lines' list according to criterion (two of them max).
   If epsilon is not -1, the second (if any) criterion is used when the difference
   of values for the first criterion is less than epsilon.
   N.B: criterion = ['energy', 'aij'] is actually the only implemented case
   with epsilon for now.
   """

   if (criterion is None):
     return

   if len(criterion) > 2:
      crit = criterion[0:2]
   else:
      crit = criterion

   if epsilon != -1:
      if crit != ['energy', 'aij']:
         raise NotImplementedError("sorting with epsilon only implemented for ['energy', 'aij']")
      def cmp(line1, line2):
         diff = line1.upper_level.energy - line2.upper_level.energy
         if abs(diff) < epsilon:
            return line1.einstein_coefficient - line2.einstein_coefficient
         return diff
      lines.sort(key=cmp_to_key(cmp))  # Inplace
      return

   for i in range(len(crit)):
      if crit[i] == 'energy':
         crit[i] = 'upper_level.energy'
      else:
         if crit[i] == 'aij':
            crit[i] = 'einstein_coefficient'

   lines.sort(key=attrgetter(*crit))  # Inplace
   return


def lines2sic():
   """
   Private to LINEDB. Do not use.
   """

   try:
      pysic.gdict.lines
      pysic.delete(pysic.gdict.lines)
   except:
      pass
   #
   leng = len(linedb.linesOut)
   pysic.define("STRUCTURE LINES",True)
   if leng == 0:
      pysic.define("INTEGER LINES%N",True)
      pysic.gdict.lines.n = leng
      return
   # Definition in reverse order so that EXA LINES% give it in the right order...
   pysic.define("CHARACTER LINES%%QNLOW*12[%d]" % (leng),True)
   pysic.define("REAL LINES%%GLOW[%d]" % (leng),True)
   pysic.define("REAL LINES%%ELOW[%d]" % (leng),True)
   #
   pysic.define("CHARACTER LINES%%QNUP*12[%d]" % (leng),True)
   pysic.define("REAL LINES%%GUP[%d]" % (leng),True)
   pysic.define("REAL LINES%%EUP[%d]" % (leng),True)
   #
   pysic.define("DOUBLE LINES%%AIJ[%d]" % (leng),True)
   pysic.define("DOUBLE LINES%%FREQUENCY[%d]" % (leng),True)
   pysic.define("REAL LINES%%UNCERTAINTY[%d]" % (leng),True)
   pysic.define("CHARACTER LINES%%SPECIES*14[%d]" % (leng),True)
   #
   pysic.define("INTEGER LINES%N",True)
   pysic.gdict.lines.n = leng
   #

   for i in range(leng):
      l = linedb.linesOut[i]
      pysic.gdict.lines.species[i] = l.species
      pysic.gdict.lines.frequency[i] = l.frequency
      pysic.gdict.lines.uncertainty[i] = l.err_frequency
      pysic.gdict.lines.aij[i] = l.einstein_coefficient
      pysic.gdict.lines.eup[i] = l.upper_level.energy
      pysic.gdict.lines.gup[i] = l.upper_level.statistical_weight
      pysic.gdict.lines.qnup[i] = l.upper_level.quantum_numbers
      pysic.gdict.lines.elow[i] = l.lower_level.energy
      pysic.gdict.lines.glow[i] = l.lower_level.statistical_weight
      pysic.gdict.lines.qnlow[i] = l.lower_level.quantum_numbers


def select_command(command,fmin,fmax,species,origin,energy,einstein,sortby,s_epsilon):
   """
   Support engine for command LINEDB\\SELECT
   This engine is called by LINEDB\\SELECT but can also be called by
   any application (e.g. WEEDS\\LFIND) which needs an identical behavior:
   - searching in the USE IN selected database,
   - using a given set of selection criteria,
   - sorting with some criteria,
   - modify the list of selected lines known by LINEDB (e.g. by LINEDB\\LIST)
   """

   linedb.lineIdx = -1  # lget() current index

   try:
      linedb.linesOut = select(command, fmin, fmax, species, origin, 'All',
        energy, einstein,sortby,s_epsilon)
   except Exception as error:
      pysic.message(pysic.seve.e,command,"{0}".format(error))
      pysic.sicerror()
      return

   # (Re)build the LINES% SIC structure
   lines2sic()

