#!/usr/bin/env python
'''
Interpolate fields from a regular CISM grid into a pre-existing MPAS grid,
using either a built-in bilinear interpolation method or using weights
generated by ESMF.
'''

import sys
import numpy as np
import netCDF4
from optparse import OptionParser
import math
from collections import OrderedDict


print "== Gathering information.  (Invoke with --help for more details. All arguments are optional)\n"
parser = OptionParser()
parser.description = "This script interpolates from a CISM grid to an MPAS grid using an ESMF weight interpolation file."
parser.add_option("-c", "--cism", dest="cismFile", help="CISM grid file to input.", default="cism.nc", metavar="FILENAME")
parser.add_option("-m", "--mpas", dest="mpasFile", help="MPAS grid file to output.", default="landice_grid.nc", metavar="FILENAME")
parser.add_option("-w", "--weight", dest="weightFile", help="ESMF weight file to input.  If not included, bilinear interpolation will be used", metavar="FILENAME")
for option in parser.option_list:
    if option.default != ("NO", "DEFAULT"):
        option.help += (" " if option.help else "") + "[default: %default]"
options, args = parser.parse_args()

print "  CISM input file:  " + options.cismFile
print "  MPAS file to be modified:  " + options.mpasFile

if options.weightFile:
    interpType = 'weight'
    print "  Interpolation will be performed using ESMF-weights method, where possible, using weights file:  " + options.weightFile
    #----------------------------
    # Get weights from file
    wfile = netCDF4.Dataset(options.weightFile, 'r')
    S = wfile.variables['S'][:]
    col = wfile.variables['col'][:]
    row = wfile.variables['row'][:]
    wfile.close()
    #----------------------------
else:
    interpType = 'bilinear'
    print "  Bilinear interpolation will be performed."

print '' # make a space in stdout before further output


#----------------------------
# Define which time levels you want to use in the two files! (0-based indexing in python)
timelev = 0
timelevout = 0
#----------------------------

#----------------------------
# Map MPAS-CISM field names - add new fields here as needed
fieldInfo = OrderedDict()
fieldInfo['thickness'] =     {'CISMname':'thk',  'scalefactor':1.0, 'CISMgrid':1, 'vertDim':False}
fieldInfo['bedTopography'] = {'CISMname':'topg', 'scalefactor':1.0, 'CISMgrid':1, 'vertDim':False}
fieldInfo['sfcMassBal'] =    {'CISMname':'acab', 'scalefactor':910.0/(3600.0*24.0*365.0), 'CISMgrid':1, 'vertDim':False}  # Assuming default CISM density
#fieldInfo['temperature'] =   {'CISMname':'temp', 'scalefactor':1.0, 'CISMgrid':1, 'vertDim':True}
fieldInfo['temperature'] =   {'CISMname':'tempstag', 'scalefactor':1.0, 'CISMgrid':1, 'vertDim':True}  # pick one or the other
fieldInfo['beta'] =          {'CISMname':'beta', 'scalefactor':1.0, 'CISMgrid':0, 'vertDim':False} # needs different mapping file...
#fieldInfo['observedSpeed'] = {'CISMname':'balvel', 'scalefactor':1.0/(365.0*24.0*3600.0), 'CISMgrid':0, 'vertDim':False} # needs different mapping file...


#----------------------------

#----------------------------
#----------------------------
# Define needed functions
#----------------------------
#----------------------------

def ESMF_interp(sourceField):
    # Interpolates from the sourceField to the destinationField using ESMF weights
  try:
    # Initialize new field to 0 - required
    destinationField = np.zeros(xCell.shape)  # fields on cells only
    sourceFieldFlat = sourceField.flatten()  # Flatten source field
    for i in range(len(row)):
      destinationField[row[i]-1] = destinationField[row[i]-1] + S[i] * sourceFieldFlat[col[i]]
  except:
     'error in ESMF_interp'
  return destinationField

#----------------------------

def BilinearInterp(Value, CISMgridType):
    # Calculate bilinear interpolation of Value field from x, y to new ValueCell field (return value)  at xCell, yCell
    # This assumes that x, y, Value are regular CISM style grids and xCell, yCell, ValueCell are 1-D unstructured MPAS style grids
  try:
    if CISMgridType == 0:
        x = x0; y = y0
    elif CISMgridType == 1:
        x = x1; y = y1
    else:
        sys.exit('Error: unknown CISM grid type specified.')
    dx = x[1] - x[0]
    dy = y[1] - y[0]
    ValueCell = np.zeros(xCell.shape)
    for i in range(len(xCell)):
       # Calculate the CISM grid cell indices (these are the lower index)
       xgrid = math.floor( (xCell[i]-x[0]) / dx )
       if xgrid >= len(x) - 1:
          xgrid = len(x) - 2
       elif xgrid < 0:
          xgrid = 0
       ygrid = math.floor( (yCell[i]-y[0]) / dy )
       if ygrid >= len(y) - 1:
          ygrid = len(y) - 2
       elif ygrid < 0:
          ygrid = 0
       #print xgrid, ygrid
       ValueCell[i] = Value[ygrid,xgrid] * (x[xgrid+1] - xCell[i]) * (y[ygrid+1] - yCell[i]) / (dx * dy) + \
                 Value[ygrid+1,xgrid] * (x[xgrid+1] - xCell[i]) * (yCell[i] - y[ygrid]) / (dx * dy) + \
                 Value[ygrid,xgrid+1] * (xCell[i] - x[xgrid]) * (y[ygrid+1] - yCell[i]) / (dx * dy) + \
                 Value[ygrid+1,xgrid+1] * (xCell[i] - x[xgrid]) * (yCell[i] - y[ygrid]) / (dx * dy) 
  except:
     'error in BilinearInterp'
  return ValueCell

#----------------------------

def interpolate_field(MPASfieldName):

    CISMfieldName = fieldInfo[MPASfieldName]['CISMname']
    if 'time' in CISMfile.variables[CISMfieldName].dimensions:
        CISMfield = CISMfile.variables[CISMfieldName][timelev,:,:]
    else:
        CISMfield = CISMfile.variables[CISMfieldName][:,:]

    print '  CISM %s min/max:'%CISMfieldName, CISMfield.min(), CISMfield.max()

    # Call the appropriate routine for actually doing the interpolation
    if interpType == 'bilinear' or fieldInfo[MPASfieldName]['CISMgrid'] == 0:  # for now, there is no capability to use a second weight file for the staggered grid - but would be easy to add later
        print "  ...Interpolating to %s using built-in bilinear method..." % MPASfieldName
        MPASfield = BilinearInterp(CISMfield, fieldInfo[MPASfieldName]['CISMgrid'])
    else:
        print "  ...Interpolating to %s using ESMF-weights method..." % MPASfieldName
        MPASfield = ESMF_interp(CISMfield)
    print '  interpolated MPAS %s min/max:'%MPASfieldName, MPASfield.min(), MPASfield.max()

    if fieldInfo[MPASfieldName]['scalefactor'] != 1.0:
        MPASfield *= fieldInfo[MPASfieldName]['scalefactor']
        print '  scaled MPAS %s min/max:'%MPASfieldName, MPASfield.min(), MPASfield.max()

    return MPASfield

    del CISMfield, MPASfield

#----------------------------

def interpolate_field_with_layers(MPASfieldName):

    CISMfieldName = fieldInfo[MPASfieldName]['CISMname']
    if 'time' in CISMfile.variables[CISMfieldName].dimensions:
        CISMfield = CISMfile.variables[CISMfieldName][timelev,:,:]
    else:
        CISMfield = CISMfile.variables[CISMfieldName][:,:]

    # create array for interpolated CISM field at all layers
    cismVerticalDimSize = CISMfield.shape[0] # vertical index is the first (since we've eliminated time already)
    mpas_grid_cism_layers = np.zeros( (cismVerticalDimSize, nCells) ) # make it the size of the CISM vertical layers, but the MPAS horizontal layers

    for z in range(cismVerticalDimSize):
        print '  CISM %s, layer %s min/max:'%(z,CISMfieldName), CISMfield[z,:,:].min(), CISMfield[z,:,:].max()
        # Call the appropriate routine for actually doing the interpolation
        if interpType == 'bilinear' or fieldInfo[MPASfieldName]['CISMgrid'] == 0:  # for now, there is no capability to use a second weight file for the staggered grid - but would be easy to add later
            print "  ...Layer %s, Interpolating to %s using built-in bilinear method..." % (z, MPASfieldName)
            mpas_grid_cism_layers[z,:] = BilinearInterp(CISMfield[z,:,:], fieldInfo[MPASfieldName]['CISMgrid'])
        else:
            print "  ...Layer %s, Interpolating to %s using ESMF-weights method..." % (z, MPASfieldName)
            mpas_grid_cism_layers[z,:] = ESMF_interp(CISMfield[z,:,:])
        print '  interpolated MPAS %s, layer %s min/max:'%(MPASfieldName, z), mpas_grid_cism_layers[z,:].min(), mpas_grid_cism_layers[z,:].max()

    if fieldInfo[MPASfieldName]['scalefactor'] != 1.0:
        mpas_grid_cism_layers *= fieldInfo[MPASfieldName]['scalefactor']
        print '  scaled MPAS %s on CISM vertical layers, min/max:'%MPASfieldName, MPASfield.min(), MPASfield.max()

    # ------------
    # Now interpolate vertically
    cism_layers = CISMfile.variables[CISMfile.variables[CISMfieldName].dimensions[1]][:]  # 2nd dimension is the vertical one - get it's name and then get the variable that has the same name
    MPASfield = vertical_interp_MPAS_grid(mpas_grid_cism_layers, cism_layers)
    print '  MPAS %s on MPAS vertical layers, min/max of all layers:'%MPASfieldName, MPASfield.min(), MPASfield.max()

    return MPASfield


#----------------------------

def vertical_interp_MPAS_grid(mpas_grid_cism_layers, cism_layers):
    destinationField = np.zeros((nCells, nVertLevels))
    for i in range(nCells):
        destinationField[i,:] = np.interp(mpasLayerCenters, cism_layers, mpas_grid_cism_layers[:,i])
    return destinationField


#----------------------------
#----------------------------




print "=================="
print 'Gathering coordinate information from CISM and MPAS files.'


# Open the input file, get needed dimensions
try:
    CISMfile = netCDF4.Dataset(options.cismFile,'r')

    # Get the CISM dimensions if they exist
    try:
      level = len(CISMfile.dimensions['level'])
    except:
      print '  Input file is missing the dimension level.  Might not be a problem.'

    try:
      stagwbndlevel = len(CISMfile.dimensions['stagwbndlevel'])
    except:
      print '  Input file is missing the dimension stagwbndlevel.  Might not be a problem.'

    # Get CISM location variables if they exist
    try:
      x1 = CISMfile.variables['x1'][:]
      dx1 = x1[1] - x1[0]
      #print 'x1 min/max/dx:', x1.min(), x1.max(), dx1
      y1 = CISMfile.variables['y1'][:]
      dy1 = y1[1] - y1[0]
      #print 'y1 min/max/dx:', y1.min(), y1.max(), dy1

      ##x1 = x1 - (x1.max()-x1.min())/2.0  # This was for some shifted CISM grid but should not be used in general.
      ##y1 = y1 - (y1.max()-y1.min())/2.0
    except:
      print '  Input file is missing x1 and/or y1.  Might not be a problem.'
    
    try:
      x0 = CISMfile.variables['x0'][:]
      #print 'x0 min/max:', x0.min(), x0.max()
      y0 = CISMfile.variables['y0'][:]
      #print 'y0 min/max:', y0.min(), y0.max()

      ##x0 = x0 - (x0.max()-x0.min())/2.0
      ##y0 = y0 - (y0.max()-y0.min())/2.0

    except:
      print '  Input file is missing x0 and/or y0.  Might not be a problem.'

except:
    sys.exit('Error: The input file specified is either missing or lacking needed dimensions/variables.')



# Open the output file, get needed dimensions & variables
try:
    MPASfile = netCDF4.Dataset(options.mpasFile,'r+')
    try:
      nVertLevels = len(MPASfile.dimensions['nVertLevels'])
    except:
      print 'Output file is missing the dimension nVertLevels.  Might not be a problem.'

    try:
      # 1d vertical fields
      layerThicknessFractions = MPASfile.variables['layerThicknessFractions'][:]
      # build up sigma levels
      mpasLayerCenters = np.zeros( (nVertLevels,) )
      mpasLayerCenters[0] = 0.5 * layerThicknessFractions[0]
      for k in range(nVertLevels)[1:]:  # skip the first level
          mpasLayerCenters[k] = mpasLayerCenters[k-1] + 0.5 * layerThicknessFractions[k-1] + 0.5 * layerThicknessFractions[k]
    except:
      print 'Output file is missing the variable layerThicknessFractions.  Might not be a problem.'

    # '2d' spatial fields on cell centers
    xCell = MPASfile.variables['xCell'][:]
    #print 'xCell min/max:', xCell.min(), xCell.max()
    yCell = MPASfile.variables['yCell'][:]
    #print 'yCell min/max:', yCell.min(), yCell.max()
    nCells = len(MPASfile.dimensions['nCells'])

except:
    sys.exit('Error: The output grid file specified is either missing or lacking needed dimensions/variables.')
print "==================\n"

# Check the overlap of the grids
print '=================='
print 'CISM File extents:'
print '  x1 min, max:    ', x1.min(), x1.max()
print '  y1 min, max:    ', y1.min(), y1.max()
print 'MPAS File extents:'
print '  xCell min, max: ', xCell.min(), xCell.max()
print '  yCell min, max: ', yCell.min(), yCell.max()
print '=================='



#----------------------------
# try each field.  If it exists in the input file, it will be copied.  If not, it will be skipped.
for MPASfieldName in fieldInfo:
  try:
    print '\n## %s ##'%MPASfieldName


    if fieldInfo[MPASfieldName]['vertDim']:
      MPASfield = interpolate_field_with_layers(MPASfieldName)
    else:
      MPASfield = interpolate_field(MPASfieldName)

    # Don't allow negative thickness.
    if MPASfieldName == 'thickness' and MPASfield.min() < 0.0:
        MPASfield[MPASfield < 0.0] = 0.0
        print '  removed negative thickness, new min/max:'%MPASfieldName, MPASfield.min(), MPASfield.max()

    # Now insert the MPAS field into the file.
    if 'Time' in MPASfile.variables[MPASfieldName].dimensions:
        MPASfile.variables[MPASfieldName][timelevout,:] = MPASfield  # Time will always be leftmost index
    else:
        MPASfile.variables[MPASfieldName][:] = MPASfield

    MPASfile.sync()  # update the file now in case we get an error later

  except:
    print '  problem with %s field (e.g. not found in input file), skipping...'%MPASfieldName



CISMfile.close()
MPASfile.close()

print '\nInterpolation completed.'
