#!/usr/bin/env python
# aeic_vertregrid.py
#
# DESCRIPTION:
# This routine vertically interpolates the AEIC aircraft emissions from
# AEIC levels onto GEOS-Chem levels (in m). Standard atmosphere altitudes
# are used for the GEOS-Chem levels (defined below). The mapping is done
# in a mass-conservative manner, e.g. the total emissions (in kg/s or 
# kg/m2/s) remain the same.
#
# HISTORY:
# 20150110 - C.Keller@seas.harvard.edu - Initial version.
# ----------------------------------------------------------------------
# Dependencies 
# ----------------------------------------------------------------------
import sys
import glob
import numpy as np
import matplotlib.dates as dates
import time
import math
from bisect import bisect_right, bisect_left
from datetime import datetime, timedelta
from netCDF4  import Dataset,  num2date, date2index, date2num

# define input and output file
infile  = "AEIC.nc"
outfile = "AEIC.47L.gen.1x1.nc"

# options:
# 1. Convert to kg/m2/s? If False, units are kept in kg/s
doArea = True 
areafile = "area.nc"

# Expand to 72 levels? If True, the AEIC data is mapped onto 72 levels and a reversed axis ('downwards'), so that the file can
# be used in GEOS-5. The original AEIC data only extends up to level 33 of the GEOS-Chem grid. Up to this level, the native and
# reduced GEOS-Chem grids are exactly the same, so we simply append 25 'fake' levels to the 47 GEOS-Chem levels defined up there.
# Emissions in all these levels will be zero anyways.
fullgrid = False 

# (reduced) GEOS-Chem levels [m] are defined here
gclevs = [ 58, 189, 320, 454, 589, 726, 864, 1004, 1146, 1290, 1436, 1584, 1759, 1988, 2249, 2517, 2792, 3074, 3439, 3896, 4375, 4879, 5413, 5980, 6585, 7237, 7943, 8846, 9936, 11021, 12086, 13134, 14170, 15198, 16222, 17243, 18727, 20836, 23020, 25307, 28654, 34024, 40166, 47135, 54834, 63053, 72180 ]

# if fullgrid, expand gclevs to 72
if fullgrid:
   for ilev in range(25):
      dum = gclevs[-1] + 100
      gclevs.append(dum)
   revert = True
else:
   revert = False

# read input file
nlevs = len(gclevs)
aeic = Dataset( infile, "r" )
hgt  = aeic.variables['height']
lon  = aeic.variables['lon']
lat  = aeic.variables['lat']
tme  = aeic.variables['time']
fB   = aeic.variables['fuelBurn']
NOx  = aeic.variables['NOx']
HC   = aeic.variables['HC']
CO   = aeic.variables['CO']

# create output file (GEOS-Chem levels)
newf = Dataset( outfile, "w" )
newf.description = "AEIC emissions. Regridded from original AEIC levels onto standard GEOS-Chem levels using routine `aeic_vertgrid.py`."
newf.history = 'Created by Christoph Keller, ' + time.ctime(time.time())

# define output dimensions
newf.createDimension('lon', len(lon[:]))
newf.createDimension('lat', len(lat[:]))
newf.createDimension('lev', nlevs)
newf.createDimension('time',None)

# get units
if doArea:
  units = "kg/m2/s"
else:
  units = "kg/s"

# define variables
lonn = newf.createVariable('lon','f4',('lon'))
lonn[:] = lon[:]
lonn.units = lon.units

latn = newf.createVariable('lat','f4',('lat'))
latn[:] = lat[:]
latn.units = lat.units

levn = newf.createVariable('lev','f4',('lev'))
levss = range(1,nlevs+1)
if revert:
   levn[:] = levss[::-1]
   levn.positive = "down"
else:
   levn[:] = levss[:]
   levn.positive = "up"

levn.long_name = "GEOS-Chem level"
levn.units = "level"

timen = newf.createVariable('time','f8',('time'))
timen.units = tme.units 
timen.calendar = 'standard'

# data variables
fBn = newf.createVariable('FUELBURN','f4',('time','lev','lat','lon'),zlib=True)
fBn.units = units
fBn.long_name = "AEIC aircraft fuel burned"

NOn = newf.createVariable('NO2','f4',('time','lev','lat','lon'),zlib=True)
NOn.units = units
NOn.long_name = "AEIC aircraft emitted NO2"

HCn = newf.createVariable('HC','f4',('time','lev','lat','lon'),zlib=True)
HCn.units = units
HCn.long_name = "AEIC aircraft emitted hydrocarbons"

COn = newf.createVariable('CO','f4',('time','lev','lat','lon'),zlib=True)
COn.units = units
COn.long_name = "AEIC aircraft emitted CO"

# initialize output arrays
FBarr = np.zeros( (12,nlevs,len(lat[:]),len(lon[:]) ) )
NOarr = np.zeros( (12,nlevs,len(lat[:]),len(lon[:]) ) )
HCarr = np.zeros( (12,nlevs,len(lat[:]),len(lon[:]) ) )
COarr = np.zeros( (12,nlevs,len(lat[:]),len(lon[:]) ) )

# eventually get area from file
if doArea:
   area = Dataset( areafile, "r" )
   am2  = area.variables['cell_area']

# walk through all GEOS-Chem levels and map aeic levels onto them
for ilev in range(len(gclevs)):

   # get upper and bottom level
   upper = gclevs[ilev]
   if ilev==0:
      bottom = 0.0
   else:
      bottom = gclevs[ilev-1]
  
   # verbose
   print '' 
   print 'now interpolating to GEOS-Chem level ' + str(ilev+1) + ': '+str(bottom)+' - '+str(upper)

   # scan through all AEIC levels
   for i in range(len(hgt[:])):
      frac = 0.0
      # get upper and bottom aeic level
      aeicu = hgt[i]
      if i==0:
         aeicb = 0.0
      else:
         aeicb = hgt[i-1]

      # check if aeic bottom is within layer 
      if ( aeicb >= bottom and aeicb < upper ): 
         # get total fraction of this aeic layer within the GC layer
         top  = min(aeicu,upper)
         frac = (top - aeicb) / (aeicu - aeicb)

      # check if aeic top is within layer
      elif ( aeicu > bottom and aeicu <= upper ):
         bot = max(aeicb,bottom)
         frac = (aeicu - bot) / (aeicu - aeicb)
      
      # check if GEOS-Chem layer is completely within this aeic layer
      elif ( aeicb < bottom and aeicu > upper ):
         frac = (upper - bottom) / (aeicu - aeicb)
 
      # add fraction to this level (only is non-zero
      if ( frac > 0.0 ):
         # verbose 
         print 'using AEIC level ' + str(i+1) + ': '+str(aeicb)+' - '+str(aeicu)+': '+str(frac)

         FBarr[:,ilev,:,:] = FBarr[:,ilev,:,:] + ( frac * fB[:,i,:,:] )
         NOarr[:,ilev,:,:] = NOarr[:,ilev,:,:] + ( frac * NOx[:,i,:,:] )
         HCarr[:,ilev,:,:] = HCarr[:,ilev,:,:] + ( frac * HC[:,i,:,:] )
         COarr[:,ilev,:,:] = COarr[:,ilev,:,:] + ( frac * CO[:,i,:,:] )

   # eventually convert to kg/m2/s
   if doArea:
      for itime in range(12):
         FBarr[itime,ilev,:,:] = FBarr[itime,ilev,:,:] / am2[:,:]
         NOarr[itime,ilev,:,:] = NOarr[itime,ilev,:,:] / am2[:,:]
         HCarr[itime,ilev,:,:] = HCarr[itime,ilev,:,:] / am2[:,:]
         COarr[itime,ilev,:,:] = COarr[itime,ilev,:,:] / am2[:,:]

# add to array
for itime in range(12):
   timen[itime] = tme[itime]
   if revert:
      fBn[itime,:,:,:] = FBarr[itime,::-1,:,:]
      NOn[itime,:,:,:] = NOarr[itime,::-1,:,:]
      HCn[itime,:,:,:] = HCarr[itime,::-1,:,:]
      COn[itime,:,:,:] = COarr[itime,::-1,:,:]
   else:
      fBn[itime,:,:,:] = FBarr[itime,:,:,:]
      NOn[itime,:,:,:] = NOarr[itime,:,:,:]
      HCn[itime,:,:,:] = HCarr[itime,:,:,:]
      COn[itime,:,:,:] = COarr[itime,:,:,:]

# close files
newf.close()
aeic.close()

#eof