#This contains some pieces of code used to derive/adapt planetary isotoplogogue
# ratios based on ARTS built-in values for Earth (which mainly come from HITRAN).
#
#Automation of calculation/conversion scheme for isotopologue ratios is
# difficult (complex), since it depends on number of atoms of each atomic
# species within a molecule, it's possible isotopes, sometimes the structure of
# the molecule (DH=HD, while DCOOH!=HCOOD), and the layout of the isotopologue
# tags.
#Hence, here we give only some pieces of code as basis for the calculation.
# Those have to be manually adapted and executed for individual species.
#
#
#Details on approach, derivation of formulas etc. in isotopratios.txt
#######################


import numpy as N
import pdb

#--------------------
#this to be done once
#--------------------
''' copy to commandline:
import isotopratios as iso
tags,d = iso.InitIR()
'''

#-------------------------------------------
#this to be repeated for individual species
#-------------------------------------------
# to do:
# (A) set name tag
# (B) uncomment/update dhe calc (incl. permut. factor)
# (C) uncomment/update dne calc
# (D) set correct mh and mn (number of H and N atoms in species)
# (E) set correct "fac" in rescaling loop (minor-isotope refactoring)

##### (A)
''' adapt name, then copy to commandline:
name='NH3'
ir,st,en = iso.getMolInfo(name,tags)
'''

##### (B)-(E)
''' adapt (B)-(E) in IRSingSpec, then copy to commandline:
reload(iso)
iso.IRSingSpec(ir,st,en,tags,d)
'''

def IRSingSpec(ir,st,en,tags,d):
  '''
  '''
  #pdb.set_trace()
#STEP (1): derive the Earth isotopic ratio that we are going to replace
  ##### (B)
  #(1b) - D/H
  dhe = d['dheh2'] #D/H from H2 (if no molecule-specific available)
  print('D/H from H2: %.3e' %dhe)

  #(1a) - D/H
  n = 2. #REPLACE permutational factor of H properly (equal the number of
         # H-atoms that can be interchanged in position, e.g. 4 for CH4)

  #in case of only one pair of main-isotope-minor-isotope isotopologue (e.g.
  # HF and DF) we can derive internal isotopic ratio from
  dhe = (1./n)*ir[-1]/ir[0] #REPLACE indices properly
  
  #in case there's more than one main-isotope-minor-isotope pair (e.g. for H2O
  # (which has H2O-16/HDO-16, H2O-18/HDO-18, H2O-17/HDO-17), HCl, ...) we take
  # the mean
  #dhe = N.mean(ir[1:3]/ir[0])
  #n=2; dhe = N.mean(N.append((1./n)*ir[3:6]/ir[0:3],N.sqrt(ir[6]/ir[0]))) #D/H from H2O
  #n=4; dhe = N.mean((1./n)*ir[2:]/ir[0:2]) #D/H from CH4
  #n=2; dhe = N.mean(N.append((1./n)*ir[3]/ir[0],N.sqrt(ir[-1]/ir[0]))) #D/H from H2CO
  #n=1; dhe = N.mean((1./n)*ir[-2:]/ir[0:2]) #HBr,HCl
  print('D/H from internal: %.3e' %dhe)

  ##### (C)
  #(1b) - 15N/14N
  # REPLACE dne calc properly
  dne = d['dnen2'] #15N/14N from N2 (if no molecule-specific available)
  print('N4/N5 from N2: %.3e' %dne)

  #(1a) - 15N/14N
  n = 2. #REPLACE permutational factor of N properly (equal the number of
         # N-atoms that can be interchanged in position, e.g. 2 for N2)

  # in case of only one pair of main-isotope-minor-isotope isotopologue we can
  # derive internal isotopic ratio from
  #miniso=-2 #common (if both H and N can be calculated)
  miniso=1 #if only N can be calculated, e.g. HNO3, N2
  dne = (1./n)*ir[miniso]/ir[0] #REPLACE indices properly

  # in case there's more than one pair we take the mean
  #dne=N.mean(ir[1:3]/ir[0]) #N2O
  print('N4/N5 from internal: %.3e' %dne)

#STEP (2): derive the rescaling factors
  ##### (D)
  #REPLACE numbers properly. if no H/N atom in molecule, then set the
  # respective m to 0.
  mh=2; mn=0

  rvh=((dhe+1)/(d['dhv']+1))**mh; rvn=1.
  rmh=((dhe+1)/(d['dhm']+1))**mh; rmn=((dne+1)/(d['dnm']+1))**mn
  rjh=((dhe+1)/(d['dhj']+1))**mh; rjn=((dne+1)/(d['dnj']+1))**mn

  #now apply the rescaling (including minor-isotope refactoring)
  dp = N.array([[d['dhv'],d['dhm'],d['dhj']],[dne,d['dnm'],d['dnj']]])
  rp = N.array([[rvh,rmh,rjh],[rvn,rmn,rjn]])
  planet = ['Venus','Mars','Jupiter']
  for i in N.arange(dp.shape[1]):
    dhp = dp[0,i]; dnp=dp[1,i]; rph = rp[0,i]; rpn = rp[1,i]
    ##### (E)
    # REPLACE fac calc properly
    #fac = N.ones(1) #when only the main isotopologue exists, e.g. C4H2, PH3
    #fac = N.append(N.ones(3),N.append(N.ones(3)*(dhp/dhe),(dhp/dhe)**2)) #H2O
    #fac = N.append(1.,N.append(N.ones(2)*(dnp/dne),N.ones(2))) #N2O
    #fac = N.append(N.ones(2),N.ones(2)*(dhp/dhe)) #CH4
    #fac = N.append(1.,N.append((dnp/dne),(dhp/dhe))) #NH3
    #fac = N.append(N.ones(2),N.append((dnp/dne),(dhp/dhe))) #HCN
    #fac = N.append(1.,(dnp/dne)) #HNO3, N2
    #fac = N.append(N.ones(2),N.ones(2)*(dhp/dhe)) #HBr,HCl
    fac = N.append(1.,(dhp/dhe)) #HI,HF,H2
    nir = fac.size
    irp = ir[:nir]*rph*rpn*fac
    print(planet[i])
    printIsoRat(irp,tags[st:,1][:nir])




def InitIR(infile='IRfromARTSBuiltin.xml'):
  '''
  Initialises basic data needed for isotopologue recalculation
  
  Parameters
  ----------
  infile : str
      name & location of file with ARTS built-in isotopologue data, from
        >isotopologue_ratiosInitFromBuiltin
        >WriteXML("ascii", isotopologue_ratios, "IRfromARTSBuiltin.xml")

  Returns
  -------
  tags : array of str (dim: (niso,3))
      the ARTS built-in isotopologue records. the 3 second-dim elements are the
       line-start symbol '@', the species&isotoplogue tag, and the
       corresponding isotopologue ratio.
  d: dict
      holds the isotopic ratios of H and N for different planets. entries as
       below:
    dhv, dhm, dhj, dnm, dnj : float
      D/H (dh) and 15N/14N (dn) isotopic ratios for [V]enus, [M]ars, and
       [J]upiter
    dheh2, dnen2 : float
      D/H and 15N/14N Earth (fixed) isotopic ratios from H2 and N2 isotopologue
       abundances
  '''
  d = {}
  #set the planetary isotopic ratios
  d['dhv'] = 1.9e-2; d['dhm'] = 8.1e-4; d['dhj'] = 2.6e-5 #D/H
  d['dnm'] = 5.7e-3; d['dnj'] = 2.25e-3

  #get ARTS built-in data
  tags=N.loadtxt(infile,comments='<',dtype=N.str)

  #inititialise D/H and 15N/14N from H2 and N2 (we don't want to do this over and over...)
  ir,st,en=getMolInfo('H2',tags,do_print=0)
  # from main isotopologue only
  #d['dheh2'] = 1./N.sqrt(ir[0])-1.
  # considering both isotopologues
  d['dheh2'] = (1./2)*ir[1]/ir[0]
  ir,st,en=getMolInfo('N2',tags,do_print=0)
  # from main isotopologue only
  #d['dnen2'] = 1./N.sqrt(ir[0])-1.
  # considering both isotopologues
  d['dnen2'] = (1./2)*ir[1]/ir[0]

  return tags, d


def getMolInfo(name,tags,do_print=1):
  '''
  Derives ARTS built-in isotopologue ratios for a specific species and their
   position&extend in the isotopologue record array

  Parameters
  ----------
  name : str
      name tag of the molecular species (= chemical formula)
  tags : array of str (dim: (niso,3))
      the ARTS built-in isotopologue records. the 3 second-dim elements are the
       line-start symbol '@', the species&isotoplogue tag, and the
       corresponding isotopologue ratio.
  do_print : boolean/int
      flag whether to print out the identified isotopologue records

  Returns
  -------
  ir : array of float (dim: miso)
      isotopologue ratios for molecular species 'name'
  st, en : int
      indices of first and last+1 elements of species 'name' in isotopologue
       record array
  '''
  st=0; found=0
  while not found:                                                  
    if (name in tags[st,1]) and (len(tags[st,1].partition('-')[0])==len(name)):
      found=1
    else:
      st+=1
  en=st
  while found:                                                      
    if (name in tags[en,1]) and (len(tags[en,1].partition('-')[0])==len(name)):
      en+=1
    else:
      found=0
  #for info we print the current info
  if do_print: print(tags[st:en])
  #for practicalities store built-in IR of the current species as numbers in
  # separate array
  ir=N.array(tags[st:en,-1],dtype=N.float)
  return ir,st,en


def printIsoRat(ir,isotags):
  '''
  Prints isotopologue record for replacement in or paste-into ARTS isotopologue
   record xml file

  Parameters
  ----------
  ir : array of float (dim: miso)
      isotopologue ratios. e.g. recalculated planetary isotopologue ratios of a
       single molecular species.
  isotags : array of str (dim: miso)
      species&isotopologue tags corresponding to the ir.
  '''
  for i,l in enumerate(ir):
    print('@ %s %.6e' %(isotags[i], l))

