# vamdc.py -- linedb implementation for VAMDC type dbs

import xml.etree.ElementTree as ET
import urllib.request, urllib.error, urllib.parse
from linedb import db
from linedb import line
from .consts import *

maxFreqMHz = 2.e6

class Vamdc(db.Db):

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

      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)
      """

      # Convert the frequencies into wavelenghts
      if fmin < 0.:
         fmin = 0.
      if fmax < 0.:
         fmax = maxFreqMHz
      #
      # Construct the URL for the database query
      vamdc_url  = self.url + "LANG=VSS2&REQUEST=doQuery&FORMAT=XSAMS&QUERY="
      vss2_query = "select *"
      #
      # Frequency criteria (RadTransFrequency in MHz)
      vss2_query += " where (RadTransFrequency >= %e and RadTransFrequency <= %e)" % (fmin,fmax)
      #
      # Einstein coefficient criterion
      if (einstein >= 0):
        vss2_query += " and RadTransProbabilityA >= %e" % einstein
      #
      # Species formula criterion
      if (len(species)>0):
        inchikeys = self.formula2inchikey(species)
        nkeys = len(inchikeys)
        if (nkeys > 0):
          for ikey in range(nkeys):
            if (ikey == 0):
              vss2_query += " and InChIKey in ('%s'" % inchikeys[ikey]
            else:
              vss2_query += ",'%s'" % inchikeys[ikey]
          vss2_query += ")"
      #
      # print vss2_query
      vss2_query = urllib.parse.quote(vss2_query)
      tap_url = vamdc_url + vss2_query
      # print tap_url
      #
      try:
        response = urllib.request.urlopen(tap_url)
      except Exception as error:
        raise Exception("Could not connect to database: %s" % error)
      #
      try:
        root = ET.parse(response).getroot()
        return self.allradiativetransition(root,energy)
      except Exception as error:
        raise Exception("Can't parse the response from the database: %s" % error)

  def part_function(self,species,origin,dbsource):
      # Returns the partition function at different temperatures
      #
      if origin.lower() != self.name:  # Case-insensitive
         raise ValueError("Got %s, but want cdms-vamdc as origin for partfunc in cdms-vamdc" % origin)
         #
      if dbsource != self.name:
         raise ValueError("Got %s, but want cdms-camdc as dbsource for partfunc in cdms-camdc" % dbsource)
         #
      self.fetch_species_dict()  # Get species dictionary if not yet done
      #
      tlist = []
      qlist = []
      for ident in list(self.moldict.keys()):
        if (self.moldict[ident]['formula'] == species):
          tlist = self.moldict[ident]['T']
          qlist = self.moldict[ident]['Q']
          #
      if (tlist==[] or qlist==[]):
        # Not found, or found but empty
        raise db.NotFoundError("No partition function found for this species.")
        #
      return tlist,qlist

  def formula2inchikey(self,formulas):
    # Return a list of InChiKey(s) matching the input list of formula(s)
    #
    self.fetch_species_dict()  # Build self.moldict if not yet done
    #
    # Build the lists of matched formulas and their InChIKeys
    lform = []
    lkeys = []
    for ident in list(self.moldict.keys()):
      if (self.moldict[ident]['formula'] in formulas):
        lform.append(self.moldict[ident]['formula'])
        lkeys.append(self.moldict[ident]['inchikey'])
    #
    # Catch unmatched species
    for formula in formulas:
      if (formula not in lform):
        raise db.NotFoundError("No such species " + formula)
    #
    # print formulas,lkeys
    #
    return list(set(lkeys))  # Remove duplicates

  def fetch_species_dict(self):
    # Fetch a minimalist dictionary of all species. Because we ask only for
    # their 'Species' table, some fields are filled (e.g. partition
    # functions), some others are not (e.g. molecular states, useful for
    # transitions).
    #
    if (hasattr(self,'moldict')):
      return
    #
    # Request results in VAMDC-XSAMS format
    vamdc_url  = self.url + "LANG=VSS2&REQUEST=doQuery&FORMAT=XSAMS&QUERY="
    vss2_query = "select Species" # All species, only species (no transitions)
    vss2_query = urllib.parse.quote(vss2_query)
    tap_url = vamdc_url + vss2_query
    #
    try:
      response = urllib.request.urlopen(tap_url)
    except Exception as error:
      raise Exception("Could not connect to database: %s" % error)
    #
    root = ET.parse(response).getroot()
    self.moldict = self.parse_species_dict(root)

  def parse_species_dict(self,root):
    # Parse the 'Species' table from the server response
    #
    ns = '{http://vamdc.org/xml/xsams/1.0}'  # Namespace for all variables...
    #
    #--- Species (Molecules) ---
    # Collect all molecules (and their possible states) in a dictionary.
    # Identifier in the "speciesID".
    #
    Species = root.find(ns+'Species')
    Molecules = Species.find(ns+'Molecules')
    #
    moldict = {}
    for Molecule in Molecules.findall(ns+'Molecule'):
      speciesId = Molecule.get('speciesID')
      MolecularChemicalSpecies = Molecule.find(ns+'MolecularChemicalSpecies')
      OrdinaryStructuralFormula = MolecularChemicalSpecies.find(ns+'OrdinaryStructuralFormula')
      Value = OrdinaryStructuralFormula.find(ns+'Value')
      formula = Value.text
      #
      InChi = MolecularChemicalSpecies.find(ns+'InChI')
      inchi = InChi.text
      InChiKey = MolecularChemicalSpecies.find(ns+'InChIKey')
      inchikey = InChiKey.text
      VAMDCSpeciesID = MolecularChemicalSpecies.find(ns+'VAMDCSpeciesID')
      vamdcspeciesid = VAMDCSpeciesID.text
      #
      # Partition function
      PartitionFunction = MolecularChemicalSpecies.find(ns+'PartitionFunction')
      if PartitionFunction is None:
        # print formula + ' has no partition function'
        Tlist = []
        Qlist = []
      else:
        T = PartitionFunction.find(ns+'T')
        Tdata = T.find(ns+'DataList')
        Tlist = [float(x) for x in Tdata.text.split()]
        Q = PartitionFunction.find(ns+'Q')
        Qdata = Q.find(ns+'DataList')
        Qlist = [float(x) for x in Qdata.text.split()]
      #
      # All molecular states
      mstates = {}
      for MolecularState in Molecule.findall(ns+'MolecularState'):
        stateID = MolecularState.get('stateID')
        MolecularStateCharacterisation = MolecularState.find(ns+'MolecularStateCharacterisation')
        StateEnergy = MolecularStateCharacterisation.find(ns+'StateEnergy')
        Value = StateEnergy.find(ns+'Value')
        units = Value.get('units')
        TotalStatisticalWeight = MolecularStateCharacterisation.find(ns+'TotalStatisticalWeight')
        if TotalStatisticalWeight is None:
          # print "No TotalStatisticalWeight for "+formula
          statweight = 0.0
        else:
          statweight = float(TotalStatisticalWeight.text)
        #
        mstates[stateID] = {}
        if (units == '1/cm'):
          mstates[stateID]['energy']     = float(Value.text)*cm_K
          mstates[stateID]['energyunit'] = 'K'
        else:
          raise Error("Energy units not supported")
        mstates[stateID]['statweight'] = statweight
        mstates[stateID]['QNs']        = ''  # ZZZ
      #
      # Save all molecule details in a dictionary
      moldict[speciesId] = {}
      moldict[speciesId]['formula']  = formula
      moldict[speciesId]['inchi']    = inchi
      moldict[speciesId]['inchikey'] = inchikey
      moldict[speciesId]['vamdcspeciesid'] = vamdcspeciesid
      moldict[speciesId]['T']        = Tlist
      moldict[speciesId]['Q']        = Qlist
      moldict[speciesId]['states']   = mstates
      #
    if (False):
      # Debug display
      print(len(list(moldict.keys())),"molecules in table. Details:")
      myformat = "%-10s: %-26s %-34s %-30s %-30s"
      print(myformat % ("speciesID","OrdinaryStructuralFormula","InChI","InChIKey","VAMDCSpeciesID"))
      for ident in list(moldict.keys()):
        print(myformat %  \
          (ident, \
          moldict[ident]['formula'], \
          moldict[ident]['inchi'], \
          moldict[ident]['inchikey'], \
          moldict[ident]['vamdcspeciesid']))
          #
    return moldict

  def allradiativetransition(self,root,energy):
    # Produce the list of transitions from server response, with the
    # enery criterion if any.
    #
    # Produce an exhaustive dictionary (all fields filled), but for the
    # the current selection only!
    mymoldict = self.parse_species_dict(root)
    #
    ns = '{http://vamdc.org/xml/xsams/1.0}'  # Namespace for all variables...
    #
    Processes = root.find(ns+'Processes')
    Radiative = Processes.find(ns+'Radiative')
    allradiativetransition = Radiative.findall(ns+'RadiativeTransition')
    # print "\nFound",len(allradiativetransition),"radiative transitions"
    #
    lines = []
    for radiativetransition in allradiativetransition:
      l = line.line()
      l.origin = self.name
      l.dbsource = self.name
      (l.species, \
       l.frequency, \
       l.err_frequency, \
       l.einstein_coefficient, \
       l.lower_level.energy, \
       l.lower_level.statistical_weight, \
       l.lower_level.quantum_numbers, \
       l.upper_level.energy, \
       l.upper_level.statistical_weight, \
       l.upper_level.quantum_numbers) \
      = self.radiativetransition2linedb(mymoldict,radiativetransition)
      if (energy > 0. and l.upper_level.energy > energy):
        continue  # cycle
      lines.append(l)
      #
    return lines

  def radiativetransition2linedb(self,moldict,RadiativeTransition):
    # Convert a RadiativeTransition to the usual linedb elements
    #
    ns = '{http://vamdc.org/xml/xsams/1.0}'  # Namespace for all variables...
    #
    SpeciesRef = RadiativeTransition.find(ns+'SpeciesRef')
    sid = SpeciesRef.text  # Species identifier
    #
    # Species name
    spec = moldict[sid]['formula']
    #
    # Frequency. ZZZ beware, there can be several frequencies (e.g. from
    # several references?)...
    EnergyWavelength = RadiativeTransition.find(ns+'EnergyWavelength')
    Frequency = EnergyWavelength.find(ns+'Frequency')
    Value = Frequency.find(ns+'Value')
    freq = float(Value.text)
    unit = Value.get('units')  # ZZZ We should check if the unit is MHz
    #
    # Frequency error
    Accuracy = Frequency.find(ns+'Accuracy')
    errfreq = float(Accuracy.text)
    #
    # Einstein coefficient
    Probability = RadiativeTransition.find(ns+'Probability')
    TransitionProbabilityA = Probability.find(ns+'TransitionProbabilityA')
    if TransitionProbabilityA is None:
      einstein = 0.0
    else:
      Value = TransitionProbabilityA.find(ns+'Value')
      einstein = float(Value.text)
    #
    # Lower and upper level characteristics
    UpperStateRef = RadiativeTransition.find(ns+'UpperStateRef')
    ulid = UpperStateRef.text
    LowerStateRef = RadiativeTransition.find(ns+'LowerStateRef')
    llid = LowerStateRef.text
    #
    # Energies
    llenergy = moldict[sid]['states'][llid]['energy']
    unit     = moldict[sid]['states'][llid]['energyunit']  # ZZZ We should check unit is K
    ulenergy = moldict[sid]['states'][ulid]['energy']
    unit     = moldict[sid]['states'][ulid]['energyunit']  # ZZZ We should check unit is K
    #
    # Statistical weights
    llsweight = moldict[sid]['states'][llid]['statweight']
    ulsweight = moldict[sid]['states'][ulid]['statweight']
    #
    # Quantum numbers
    # ZZZ We have to re-encode the v, J, F quantum numbers found in XSAMS into the
    # former format. See details in https://www.astro.uni-koeln.de/node/477#qnfmt
    llqns = ''
    ulqns = ''
    #
    return (spec,freq,errfreq,einstein,llenergy,llsweight,llqns,ulenergy,ulsweight,ulqns)
