# This is a module file for handling GRIB data I/O and metadata. This module
# file relies heavily on pygrid for reading in data from GRIB edition 1 
# and 2 files. 

# Logan Karsten
# National Center for Atmospheric Research 
# Research Applications Laboratory
# karsten@ucar.edu
# 303-497-2693

#import pygrib
import numpy as np
import os
import subprocess
from netCDF4 import Dataset

class gribMeta:
    def __init__(self):
        self.metaInitialized = False
        self.nx = -9999
        self.ny = -9999
        self.foundNdv = False
        self.gribNdv = -9999
        self.gribEdition = -9999
        self.dx = -9999
        self.dy = -9999
        self.projection = 'MISSING'
        self.latsFound = False
        self.gridCenterLats = -9999
        self.lonsFound = False
        self.gridCenterLons = -9999
        self.lon0 = -9999
        self.lat0 = -9999
        self.lon1 = -9999
        self.lat1 = -9999
        self.lonV = -9999
        self.eRadius = -9999
        self.sphericalRadius = -9999
        self.majorRadius = -9999
        self.minorRadius = -9999
        self.latD = -9999
        self.lonOrigin = -9999
    def readMeta(self,statusMeta,gribFile,id):
        """
        Generic routine for opening GRIB files, and looking for some expected 
        metadata. This includes information on grid lat/lon values, number 
        of rows and columns, grid resolution, and projection information. 
        It's important to read this in every time to protect against unexpected
        grid changes that need to be accounted for on the forcing end. 
        """
        # First, we will creat a NetCDF dump from the GRIB file to the temporary scratch
        # directory. From there, we will read in some metadata from the NetCDF file.
        tmpNcPath = statusMeta.tmpDir + "/tmpGRIB.nc"

        # If temporary path exists, remove it from old workflow instance.
        if os.path.isfile(tmpNcPath):
            try:
                os.remove(tmpNcPath)
            except:
                statusMeta.statusMsg = "Unable to remove old NetCDF file: " + tmpNcPath
                raise Exception()
        
        # Create a NetCDF file.
        if statusMeta.ncepWcoss == 1:
            cmd = "$WGRIB2 " + gribFile + " -match :\"" + id + "\": -netcdf " + tmpNcPath
        else:
            cmd = "wgrib2 " + gribFile + " -match :\"" + id + "\": -netcdf " + tmpNcPath

        try:
            # set up GRIB2TABLE if needed:
            if not os.environ.get('GRIB2TABLE'):
                g2path = os.path.join(statusMeta.tmpDir, "grib2.tbl")
                with open(g2path, 'wt') as g2t:
                    g2t.write(
                            "209:1:0:0:161:1:6:30:MultiSensorQPE01H:"
                            "Multi-sensor estimated precipitation accumulation 1-hour:mm\n"
                            "209:1:0:0:161:1:6:37:MultiSensorQPE01H:"
                            "Multi-sensor estimated precipitation accumulation 1-hour:mm\n"
                    )
                os.environ['GRIB2TABLE'] = g2path
 
            cmdOutput = subprocess.Popen([cmd],stdout=subprocess.PIPE,stderr=subprocess.PIPE,shell=True)
            out, err = cmdOutput.communicate()
            exitcode = cmdOutput.returncode
        except:
            statusMeta.statusMsg = "Unable to create temporary NetCDF file from: " + gribFile
            raise Exception()

        if exitcode != 0:
            statusMeta.statusMsg = "wgrib2 failed on: " + gribFile
            raise Exception()

        # Open the file.
        try:
            idTmp = Dataset(tmpNcPath)
        except:
            statusMeta.statusMsg = "Unable to open temporary NetCDF file: " + tmpNcPath
            raise Exception()

        if not 'latitude' in list(idTmp.variables):
            statusMeta.statusMsg = "Expected latitude not found in: " + tmpNcPath
            raise Exception()

        if not 'longitude' in list(idTmp.variables):
            statusMeta.statusMsg = "Expected longitude not found in: " + tmpNcPath
            raise Exception()

        if len(idTmp.variables['latitude'].shape) == 1:
            self.nx = idTmp.variables['longitude'].shape[0]
            self.ny = idTmp.variables['latitude'].shape[0]
        else:
            self.nx = idTmp.variables['latitude'].shape[1]
            self.ny = idTmp.variables['latitude'].shape[0]

        latCoords = np.empty([self.ny,self.nx],np.float32)
        lonCoords = np.empty([self.ny,self.nx],np.float32)

        if len(idTmp.variables['latitude'].shape) == 1:
            for i in range(0,self.nx):
                latCoords[:,i] = idTmp.variables['latitude'][:]
            for j in range(0,self.ny):
                lonCoords[j,:] = idTmp.variables['longitude'][:]
        else:
            #if not 'x' in list(idTmp.variables):
            #    statusMeta.errMsg = "ERROR: Expected x variable not found in: " + tmpNcPath
            #    raise Exception()
            #    #if not 'y' in list(idTmp.variables):
            #    statusMeta.errMsg = "ERROR: Expected y variable not found in: " + tmpNcPath
            #    raise Exception()
            if len(idTmp.variables['latitude'].shape) != 2:
                statusMeta.statusMsg = "Unexpected latitude size in: " + tmpNcPath
                raise Exception()
            if len(idTmp.variables['longitude'].shape) != 2:
                statusMeta.statusMsg = "Unexpected longitude size in: " + tmpNcPath
                raise Exception()
            latCoords[:,:] = idTmp.variables['latitude'][:,:]
            lonCoords[:,:] = idTmp.variables['longitude'][:,:]

        # Pull grid-center lat/lon values
        try:
            self.gridCenterLats = latCoords
            self.latsFound = True
        except:
            statusMeta.statusMsg = "Unable to extract grid center point latitude values from file: " + gribFile
            raise Exception()
            
        try:
            self.gridCenterLons = lonCoords
            self.lonsFound = True
        except:
            statusMeta.statusMsg = "Unable to extract grid center point longitude values from file: " + gribFile
            raise Exception()
            
        # Change the status of the ojbect to reflect metadata being read in correctly.
        self.metaInitialized = True

        # Close and remove the temporary NetCDF file
        try:
            idTmp.close()
        except:
            statusMeta.statusMsg = "Unable to close NetCDF file: " + tmpNcPath
            raise Exception()

        try:
            os.remove(tmpNcPath)
        except:
            statusMeta.statusMsg = "Unable to remove NetCDF file: " + tmpNcPath
            raise Exception()
   
def getGribVariable(statusMeta,gribFile,gribMeta,id,varName,levelType,lev,expectedUnits,dataArray):
    """
    Function to read in GRIB data given a file, variable name, and number of 
    vertical levels. 
    """
    # First, we will creat a NetCDF dump from the GRIB file to the temporary scratch 
    # directory. From there, we will read in some metadata from the NetCDF file.
    tmpNcPath = statusMeta.tmpDir + "/tmpGRIB.nc"

    # If temporary path exists, remove it from old workflow instance.
    if os.path.isfile(tmpNcPath):
        try:
            os.remove(tmpNcPath)
        except:
            statusMeta.statusMsg = "Unable to remove old NetCDF file: " + tmpNcPath
            raise Exception()


    # Create a NetCDF file.
    if statusMeta.ncepWcoss == 1:
        cmd = "$WGRIB2 " + gribFile + " -match :\"" + id + "\": -netcdf " + tmpNcPath
    else:
        cmd = "wgrib2 " + gribFile + " -match :\"" + id + "\": -netcdf " + tmpNcPath

    try:
        # set up GRIB2TABLE if needed:
        if not os.environ.get('GRIB2TABLE'):
           g2path = os.path.join(statusMeta.tmpDir, "grib2.tbl")
           with open(g2path, 'wt') as g2t:
                    g2t.write(
                            "209:1:0:0:161:1:6:30:MultiSensorQPE01H:"
                            "Multi-sensor estimated precipitation accumulation 1-hour:mm\n"
                            "209:1:0:0:161:1:6:37:MultiSensorQPE01H:"
                            "Multi-sensor estimated precipitation accumulation 1-hour:mm\n"
                    )
           os.environ['GRIB2TABLE'] = g2path

        cmdOutput = subprocess.Popen([cmd],stdout=subprocess.PIPE,stderr=subprocess.PIPE,shell=True)
        out, err = cmdOutput.communicate()
        exitcode = cmdOutput.returncode
    except:
        statusMeta.statusMsg = "Unable to create temporary NetCDF file from: " + gribFile
        raise Exception()

    if exitcode != 0:
        statusMeta.statusMsg = "wgrib2 failed on: " + gribFile
        raise Exception()

    # Open the file. 
    try:
        idTmp = Dataset(tmpNcPath)
    except:
        statusMeta.statusMsg = "Unable to open temporary NetCDF file: " + tmpNcPath
        raise Exception()

    # Check to make sure expected NetCDF variable is present in file.
    if varName not in list(idTmp.variables):
        statusMeta.statusMsg = "Unable to locate expected variable: " + varName + " in: " + tmpNcPath
        raise Exception()

    # Sanity checking to make sure variable shape conforms to expected array size. 
    if dataArray.shape[0] != idTmp.variables[varName].shape[1]:
        statusMeta.statusMsg = "Expected NY of: " + str(dataArray.shape[0]) + \
                            " Receivd NY of: " + str(idTmp.variables[varName].shape[1]) + \
                            " from file: " + gribFile
    if dataArray.shape[1] != idTmp.variables[varName].shape[2]:
        statusMeta.statusMsg = "Expected NY of: " + str(dataArray.shape[1]) + \
                            " Receivd NY of: " + str(idTmp.variables[varName].shape[2]) + \
                            " from file: " + gribFile

    # Extract data values from file.
    try:
        dataArray[:,:] = idTmp.variables[varName][0,:,:]
    except:
        statusMeta.statusMsg = "Unable to pull data for variable: " + varName + \
                            " from file: " + gribFile
        raise Exception()

    # Going to make a crude assumption that the fill value in the NetCDF variable here is global
    # to all variables in the file. It's a bit dangerous.
    if not gribMeta.foundNdv:
        gribMeta.gribNdv = idTmp.variables[varName]._FillValue

    # Perform a sanity check here to make sure the data is not an entire grid of
    # all missing values. 
    validInd = np.where(dataArray != idTmp.variables[varName]._FillValue)
    if len(validInd[0]) == 0:
        statusMeta.statusMsg = "No valid points found for variable: " + varName + \
                            " from file: " + gribFile
        raise Exception()
        
    # Close and remove the temporary NetCDF file
    try:
        idTmp.close()
    except:
        statusMeta.statusMsg = "Unable to close NetCDF file: " + tmpNcPath
        raise Exception()

    try:
        os.remove(tmpNcPath)
    except:
        statusMeta.statusMsg = "Unable to remove NetCDF file: " + tmpNcPath
        raise Exception()

def gribOneToGribTwo(statusMeta,gribFileIn,gribFileOut):
    """
    Generic function to utilize the cnvgrib command line utility to
    convert GRIB1 files to GRIB2
    """
    # Double check to make sure input file exists.
    if not os.path.isfile(gribFileIn):
        statusMeta.statusMsg = "Input GRIB2 file: " + gribFileIn + " not found."
        raise Exception()

    # Run cnvgrib command to convert the file.
    cmd = "cnvgrib -g12 " + gribFileIn + " " + gribFileOut
    if statusMeta.ncepWcoss == 1:
        cmd = "$CNVGRIB -g12 " + gribFileIn + " " + gribFileOut

    try:
        subprocess.call(cmd,shell=True)
    except:
        statusMeta.statusMsg = "Unable to run cnvgrib on: " + gribFileIn
        raise Exception()

    # Double check to make sure expected output file exists.
    if not os.path.isfile(gribFileOut):
        statusMeta.statusMsg = "Output GRIB2 file: " + gribFileOut + " not found."
        raise Exception()