from netCDF4 import Dataset
import numpy as np

def ncvardef ( nc, varinp, dims, data=None ):
  """ Create netCDF variable object

  Version
    22NOV24 AD Add 'lin', 'off', 'qad', 'sca'
    20NOV24 AD Add 'co2'
    30JUL24 AD Allow for scalar data, and '_res', '_csq'
    19APR24 AD Original

  Parameters
    nc     obj  : existing netCDF object 
    varinp str  : input variable short name, eg 'lat', 'tem_sd'
    dims   (:)  : either a single label or tuple of labels of array dimensions
                  empty tuple () indicates a scalar
    data  [dims]: (optional) data to be written into created object

  Returns
    If 'data' parameter specified
      None
    Else
      netCDF data object 

  Usage
    This can be used to define and load data directly, eg with tem_data
    as a 2D array of size (npix,nlev) 

      ncvardef ( nc, 'tem', (npix,nlev), tem_data ) 

    Alternatively it can be used just to define the netCDF variable, 
    with the data loaded subsequently

      ncv_tem = ncvardef ( nc, 'tem', (npix,nlev) ) 
      ncv_tem[:,:] = tem_data     
    (Note the explicit use of '[:,:]' to avoid refining ncv_tem as a simple
     2D array matching tem_data)

  """
  var = varinp
  useext = False
  for ext in ( '_sd', '_csq', '_dfs', '_inf', '_res' ):
    if ext in varinp:
      ius = varinp.find(ext)   # look for underscore in varinp
      var = varinp[0:ius]
      useext = ext
      break
    
  vardic = {}
  vardic['aerosol'] = { 'long_name' : 'aerosol extinction',
                            'units' : 'per_km' }
  vardic['c2h6'] = {    'long_name' : 'ethane',
                    'standard_name' : 'mole_fraction_of_ethane_in_air',
                            'units' : '1e-6' }
  vardic['ccl4'] = {    'long_name' : 'calcium tetrachloride',
                    'standard_name' : 'mole_fraction_of_calcium_tetrachloride_in_air',
                            'units' : '1e-6' } 
  vardic['ch4'] = {     'long_name' : 'methane',
                    'standard_name' : 'mole_fraction_of_methane_in_air',
                            'units' : '1e-6' } 
  vardic['chisq'] = {   'long_name' : 'average number of iterations for converged retrievals' } 
  vardic['cld_flg'] = { 'long_name' : 'cloud flag',
                             'type' : np.int8 }
  vardic['clono2'] = {  'long_name' : 'chlorine nitrate',
                    'standard_name' : 'mole_fraction_of_chlorine_nitrate_in_air',
                            'units' : '1e-6' }
  vardic['co'] = {      'long_name' : 'carbon monoxide',
                    'standard_name' : 'mole_fraction_of_carbon_monoxide_in_air',
                            'units' : '1e-6' }
  vardic['co2'] = {     'long_name' : 'carbon dioxide',
                    'standard_name' : 'mole_fraction_of_carbon_dioxide_in_air',
                            'units' : '1e-6' } 
  vardic['cof2'] = {    'long_name' : 'carbonyl fluoride',
                    'standard_name' : 'mole_fraction_of_carbonyl_flouride_in_air',
                            'units' : '1e-6' }
  vardic['f11'] = {     'long_name' : 'cfc11',
                    'standard_name' : 'mole_fraction_of_cfc11_in_air',
                            'units' : '1e-6' }
  vardic['f12'] = {     'long_name' : 'cfc12', 
                    'standard_name' : 'mole_fraction_of_cfc12_in_air',
                            'units' : '1e-6' }
  vardic['f14'] = {     'long_name' : 'carbon tetrafluoride', 
                    'standard_name' : 'mole_fraction_of_carbon_tetrafluoride_in_air',
                            'units' : '1e-6' }
  vardic['f22'] = {     'long_name' : 'hcfc22',
                    'standard_name' : 'mole_fraction_of_hcfc22_in_air',
                            'units' : '1e-6' }
  vardic['h2o'] = {     'long_name' : 'water vapor',
                    'standard_name' : 'mole_fraction_of_water_vapor_in_air',
                            'units' : '1e-6' }
  vardic['hcn'] = {     'long_name' : 'hydrogen cyanide',
                    'standard_name' : 'mole_fraction_of_hydrogen_cyanide_in_air',
                            'units' : '1e-6' }
  vardic['hgt'] = {     'long_name' : 'retrieved height',
                    'standard_name' : 'height',
                            'units' : 'km' }
  vardic['hgt_nom'] = { 'long_name' : 'nominal tangent height',
                            'units' : 'km' }
  vardic['hms'] = {     'long_name' : 'time as hhmmss (ut)',
                            'units' : 'hhmmss',
                             'type' : 'i4' }
  vardic['hno3'] = {    'long_name' : 'nitric acid',
                    'standard_name' : 'mole_fraction_of_nitric_acid_in_air',
                            'units' : '1e-6' }
  vardic['hocl'] = {    'long_name' : 'hypochlorous acid',
                    'standard_name' : 'mole_fraction_of_hypochlorous_acid_in_air',
                            'units' : '1e-6' }
  vardic['jsec'] = {    'long_name' : 'seconds since 2000-01-01 00:00:00',
                    'standard_name' : 'time',
                            'units' : 'seconds',
                             'type' : np.int32 }
  vardic['lat'] = {     'long_name' : 'latitude of retrieved profile center',
                    'standard_name' : 'latitude',
                            'units' : 'degree_north',
                      'valid_range' : [ -90, 90 ] }
  vardic['lon'] = {     'long_name' : 'latitude of retrieved profile center',
                    'standard_name' : 'latitude',
                            'units' : 'degree_east',
                      'valid_range' : [ -180, 180 ] }
  vardic['lin'] = {     'long_name' : 'spectral linear correction' }
  vardic['lst'] = {     'long_name' : 'local solar time', 
                            'units' : 'hours since 00:00:00 local solar midnight',
                      'valid_range' : [ 0, 24 ] }
  vardic['n2o'] = {     'long_name' : 'nitrous oxide',
                    'standard_name' : 'mole_fraction_of_nitrous_oxide_in_air',
                            'units' : '1e-6' }
  vardic['n2o5'] = {    'long_name' : 'dinitrogen pentoxide',
                    'standard_name' : 'mole_fraction_of_dinitrogen_pentoxide_in_air',
                            'units' : '1e-6' }
  vardic['no2'] = {     'long_name' : 'nitrogen dioxide', 
                    'standard_name' : 'mole_fraction_of_nitrogen_dioxide_in_air',
                            'units' : '1e-6' }
  vardic['n_iter'] = {  'long_name' : 'average number of iterations for converged retrievals' } 
  vardic['n_notcnv'] = {'long_name' : 'number of sequential retrievals not converged', 
                             'type' : 'i' }
  vardic['n_seqrtv'] = {'long_name' : 'number of sequential retrievals attempted', 
                             'type' : 'i' }
  vardic['o3'] = {      'long_name' : 'ozone',
                    'standard_name' : 'mole_fraction_of_ozone_in_air',
                            'units' : '1e-6' }
  vardic['ocs'] = {     'long_name' : 'carbonyl sulfide',
                    'standard_name' : 'mole_fraction_of_carbonyl_sulfide_in_air',
                            'units' : '1e-6' }
  vardic['off'] = {     'long_name' : 'radiometric offset',
                            'units' : 'radiance units' }
  vardic['orb'] = {     'long_name' : 'orbit number', 
                             'type' : 'i' }
  vardic['poi'] = {     'long_name' : 'pointing correction', 
                            'units' : 'km' }
  vardic['pre'] = {     'long_name' : 'pressure', 
                    'standard_name' : 'air_pressure',
                            'units' : 'hPa' }
  vardic['qad'] = {     'long_name' : 'spectral quadratic correction' }
  vardic['refpre'] = {  'long_name' : 'reference level pressure', 
                    'standard_name' : 'air_pressure',
                            'units' : 'hPa' }
  vardic['sca'] = {     'long_name' : 'radiometric scale factor correction' }
  vardic['sf6'] = {     'long_name' : 'sulfur hexafluoride',
                    'standard_name' : 'mole_fraction_of_sulfur_hexafluoride_in_air',
                            'units' : '1e-6' }
  vardic['sza'] = {     'long_name' : 'solar zenith angle',
                    'standard_name' : 'solar_zenith_angle',
                            'units' : 'degrees',
                      'valid_range' : [ 0, 180 ] }
  vardic['tem'] = {     'long_name' : 'temperature',
                    'standard_name' : 'air_temperature',
                            'units' : 'K' }
  vardic['ymd'] = {     'long_name' : 'date as yyyymmdd',
                            'units' : 'yyyymmdd',
                             'type' : 'i4' }

  if var in vardic:               # check that input var is in recognised list
    vardat = vardic[var]
    if 'type' in vardat:
      vartyp = vardat['type'] 
    else:                         # default data type is 4-byte floating number
      vartyp = 'f4'
    varobj = nc.createVariable(varinp,vartyp,dims)

    for key in vardat:            # Add data attributes
      if key == 'type': continue  # Exclude 'type' - not an attribute
      if useext == '_sd':  
        if key == 'long_name' or key == 'standard_name': vardat[key] += ' sd'
      elif useext == '_csq':
        if key == 'long_name' or key == 'standard_name': vardat[key] += ' chi-squared relative to a priori'
        if key == 'units': continue    # ignore 'units' key for dfs, which is dimensionless
      elif useext == '_dfs':
        if key == 'long_name' or key == 'standard_name': vardat[key] += ' degrees of freedom for signal'
        if key == 'units': continue    # ignore 'units' key for dfs, which is dimensionless
      elif useext == '_inf':
        if key == 'long_name' or key == 'standard_name': vardat[key] += ' information content'
        if key == 'units': vardat[key] = 'bits'
      elif useext == '_res':
        if key == 'long_name' or key == 'standard_name': vardat[key] += ' vertical resolutoin scale factor'
        if key == 'units': continue    # ignore 'units' key for res, which is dimensionless
      setattr ( varobj, key, vardat[key] )
  else:
    print('F-ncvardef: unrecognised var=' + var)
    exit()
  if data is not None: 
    if len(dims) == 0:              # scalar
      varobj[:] = data
    elif isinstance(dims,str):
      varobj[:] = data
    else:
      varobj[:,:] = data
  return varobj  

