#!/usr/bin/env python 
"""
Created on Tue Jun 25 19:59:49 2013

@author: Hessel Winsemius
"""
import netCDF4 as nc4
import numpy as np
import os
#import pdb
import datetime
import osgeo.osr as osr
import logging
import logging.handlers
import sys

def setlogger(logfilename, logReference):
    """
    Set-up the logging system. Exit if this fails
    """
    try:
        #create logger
        logger = logging.getLogger(logReference)
        logger.setLevel(logging.DEBUG)
        ch = logging.handlers.RotatingFileHandler(logfilename,maxBytes=10*1024*1024, backupCount=5)
        console = logging.StreamHandler()
        console.setLevel(logging.DEBUG)
        ch.setLevel(logging.DEBUG)
        #create formatter
        formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
        #add formatter to ch
        ch.setFormatter(formatter)
        console.setFormatter(formatter)
        #add ch to logger
        logger.addHandler(ch)
        logger.addHandler(console)
        logger.debug("File logging to " + logfilename)
        return logger, ch
    except IOError:
        print "ERROR: Failed to initialize logger with logfile: " + logfilename
        sys.exit(2)

def recursive_glob(rootdir='.', suffix=''):
    """
    Prepares a list of files in location rootdir, with a suffix
    input:
        rootdir:    string, path-name
        suffix:     suffix of required files
    output:
        fileList:   list-strings, file paths and names
    """

    fileList = [os.path.join(rootdir, filename)
            for rootdir, dirnames, filenames in os.walk(rootdir)
            for filename in filenames if filename.endswith(suffix)]
    return fileList

def prepare_nc(nc_src, FileOut, datetimeObj, x, y, lon_vals, lat_vals, srs, epsgStr, logger):
    """
    Prepares a target NetCDF, following the attributes of a source NetCDF file, but using the
    x-axis, y-axis of a certain projection. If the projection is lat-lon, then only a lat-lon axis
    is prepared. Otherwise an x, y axis is prepared, a projection variable is defined, and grids for
    latitude and longitude values are written to the NetCDF file. In principle, files are generated per year
    starting at the 1st of January.
    input:
    ::

        nc_src:         NetCDF object of source file
        FileOut:        string, referring to target NetCDF file location
        datetimeObj:    list of datetime objects, to be written to the time variable of the target NetCDF
        x:              Numpy array, containing x-axis of target grid in projection, defined by projStr
        y:              Numpy array, containing y-axis of target grid in projection, defined by projStr
        lon_vals:       Longitude values of pixels in the projected grid
        lat_vals:       Latitude values of pixels in the projected grid
        logger:         logger object

    output:
    ::

        No output from this function. The result is a prepared NetCDF file at FileOut

    """
    # if projStr.lower() == 'epsg:4326':
    if srs.IsProjected() == 0:
        logger.info('Found lat-lon coordinate system, preparing lat, lon axis')
        x_dim      = 'lon';             y_dim     = 'lat'
        x_name     = 'longitude';       y_name = 'latitude'
        x_longname = 'Longitude values';y_longname = 'Latitude values'
        x_unit     = 'degrees_east';       y_unit     = 'degrees_north'
        gridmap    = 'latitude_longitude'
    else:
        logger.info('Found a Cartesian projection coordinate system, preparing x, y axis')
        x_dim = 'x';                                    y_dim = 'y'
        x_name = 'projection_x_coordinate';             y_name = 'projection_x_coordinate'
        x_longname = 'x coordinate of projection';      y_longname = 'y coordinate of projection'
        x_unit     = 'm';                               y_unit     = 'm'
        gridmap    = 'Cartesian'

    logger.info('Preparing ' + FileOut)
    nc_trg = nc4.Dataset(FileOut,'w') # format='NETCDF3_CLASSIC'
    # Create dimensions
    nc_trg.createDimension("time", 0) #NrOfDays*8
    nc_trg.createDimension(y_dim, len(y))
    nc_trg.createDimension(x_dim, len(x))
    # create axes

    DateHour = nc_trg.createVariable('time','f8',('time',))
    try:
        DateHour.units = nc_src.variables['time'].units
    except:
        DateHour.units = 'Days since 1970-01-01 00:00:00'
    try:
        DateHour.calendar = nc_src.variables['time'].calendar
    except:
        DateHour.calendar = 'gregorian'
    try:
        DateHour.standard_name = nc_src.variables['time'].standard_name
    except:
        DateHour.standard_name = 'time'
    try:
        DateHour.long_name = nc_src.variables['time'].long_name
    except:
        DateHour.long_name = 'time'
    DateHour[:] = nc4.date2num(datetimeObj,units=nc_src.variables['time'].units,calendar=DateHour.calendar)
    y_var = nc_trg.createVariable(y_dim,'f4',(y_dim,))
    y_var.standard_name = y_name
    y_var.long_name = y_longname
    y_var.units = y_unit
    x_var = nc_trg.createVariable(x_dim,'f4',(x_dim,))
    x_var.standard_name = x_name
    x_var.long_name = x_longname
    x_var.units = x_unit
    y_var[:] = y
    x_var[:] = x

    # if projStr.lower() != 'epsg:4326':
    if srs.IsProjected() == 1:
        lat = nc_trg.createVariable('lat','f4',(y_dim,x_dim,))
        lat.standard_name = 'latitude'
        lat.long_name = 'latitude coordinate'
        lat.units = 'degrees_north'
        lat.coordinates = 'lat lon'
        lat.grid_mapping = 'wgs84'
        #lat._CoordinateAxisType = "Lat"
        lat[:,:] = lat_vals
        lon = nc_trg.createVariable('lon','f4',(y_dim,x_dim,))
        lon.standard_name = 'longitude'
        lon.long_name = 'longitude coordinate'
        lon.units = 'degrees_east'
        lon.coordinates = 'lat lon'
        lon.grid_mapping = 'wgs84'
        #lon._CoordinateAxisType = "Lon"
        lon[:,:] = lon_vals

    # Set attributes
    # Change some of the attributes, add some
    nc_trg.Conventions     = 'CF-1.6'
    nc_trg.publisher_name = 'Hessel Winsemius'
    nc_trg.publisher_email = 'hessel.winsemius@deltares.nl'
    nc_trg.publisher_url = 'http://www.deltares.nl'
    nc_trg.title           = nc_src.title
    nc_trg.license = 'These data are licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License.'
    try:
        nc_trg.institution     = nc_src.institution
    except:
        nc_trg.institution     = 'unknown'
    nc_trg.history = '{:s} created by yearly stacking by stack_yearly.py'.format(datetime.datetime.utcnow().strftime('%Y-%m-%dT%H:%M:%SZ'))
    nc_trg.disclaimer      = 'The availability and quality of these data is in no way guaranteed by Deltares'
    nc_trg.naming_authority = 'esa.int'
    nc_trg.id = 'gfo_release_{:s}'.format(datetime.datetime.utcnow().strftime('%Y%m%d'))
    nc_trg.time_coverage_resolution = '1 day'
    nc_trg.geospatial_lat_min = lat_vals.min()
    nc_trg.geospatial_lat_max = lat_vals.max()
    nc_trg.geospatial_lon_min = lon_vals.min()
    nc_trg.geospatial_lon_max = lon_vals.max()
    nc_trg.geospatial_lat_units = 'degrees_north'
    nc_trg.geospatial_lon_units = 'degrees_east'
    nc_trg.geospatial_lat_resolution = 0.009
    nc_trg.geospatial_lon_resolution = 0.009
    nc_trg.creator_name = 'esa'
    nc_trg.creator_url = 'http://www.esa.int'
    nc_trg.project = 'Earth2Observe'
    nc_trg.processing_level = 'preliminary'
    nc_trg.references = 'http://www.hydrol-earth-syst-sci.net/17/651/2013/hess-17-651-2013.html'
    nc_trg.keywords = 'EnviSAT, ASAR, flood, probability, quality, water'
    nc_trg.keywords_vocabulary = 'NASA/GCMD Earth Science Keywords. Version 6.0'
    nc_trg.date_created = datetime.datetime.utcnow().strftime('%Y-%m-%dT%H:%M:%SZ')
    nc_trg.date_modified = datetime.datetime.utcnow().strftime('%Y-%m-%dT%H:%M:%SZ')
    nc_trg.date_issued = datetime.datetime.utcnow().strftime('%Y-%m-%dT%H:%M:%SZ')
    
    # write projection info to file
    projection= nc_trg.createVariable('projection','c')
    projection.long_name = 'projection'
    projection.EPSG_code = epsgStr
    projection.proj4_params = srs.ExportToProj4()
    projection.grid_mapping_name = gridmap
    wgs84 = nc_trg.createVariable('wgs84','c')
    wgs84.long_name = 'wgs84'
    wgs84.EPSG_code = 'EPSG:4326'
    wgs84.proj4_params = '+proj=longlat +ellps=WGS84 +datum=WGS84 +no_defs'
    wgs84.grid_mapping_name = 'latitude_longitude'
    nc_trg.sync()
    nc_trg.close()


srcFolder = r'd:\scratch\GFO'
trgFolder = r'd:\scratch\GFO_out'
srcFolder = r'/envisat_archive/data/4-floodmap'
trgFolder = r'/envisat_archive/netCDF_stacks'
trgbetween = '1000m' 

tileNr_X = range(-180,180,10) # range(-180,180,10)
tileNr_Y = range(-90,90,10)# range(-90,90,10)

proj = 'EPSG:4326'
# make proj string and epsg code
srs = osr.SpatialReference()
srs.ImportFromEPSG(int(proj[5:]))
projStr = srs.ExportToProj4()
epsgStr = proj

logger, ch = setlogger('stack_GFA.log','GFA')
for tile_X in tileNr_X:
    for tile_Y in tileNr_Y:
        # prepare suffix string
        if tile_X <= 0:
            easting = str('W%03.f') % (-tile_X)
        else:
            easting = str('E%03.f') % tile_X
        if tile_Y <= 0:
            northing = str('S%03.f') % (-tile_Y)
        else:
            northing = str('N%03.f') % tile_Y

        suffix = str(easting + northing + '.nc')
        # make a list of files belonging to this tile from all processed files, may take a while for 10 years
        tileFileList = recursive_glob(rootdir=srcFolder, suffix=suffix)
        tileFileList.sort()
        #pdb.set_trace()
        for tileFile in tileFileList:
            logger.info('Treating tile ' + tileFile)
            a = nc4.Dataset(tileFile)
            time = a.variables['time']
            qual = a.variables['qualityIndicator']
            prob = a.variables['waterProbability']
            # pdb.set_trace()
            timeObj = nc4.num2date(time[:], units=time.units, calendar=time.calendar)[0]
            # check the year of the tile and see if a file already exists
            year = timeObj.year
            trgFile = os.path.join(trgFolder, suffix[0:-3], str(year) + '_' + trgbetween + '_' + suffix)
            if not(os.path.isdir(os.path.split(trgFile)[0])):
                os.makedirs(os.path.split(trgFile)[0])
            if not(os.path.isfile(trgFile)):
                x = a.variables['lon'][:]
                y = a.variables['lat'][:]
                
                # prepare a new file containing the year as timeObjs
                curTime = datetime.datetime(year, 1, 1, 0, 0)
                allDates = []
                while curTime.year == year:
                    allDates.append(curTime)
                    curTime += datetime.timedelta(1)
                # print 'Preparing file ' + trgFile
                prepare_nc(a, trgFile, allDates, x, y, None, None, srs, epsgStr, logger)
                # add start time
                nc_trg = nc4.Dataset(trgFile, 'a')
                nc_trg.time_coverage_start = timeObj.strftime('%Y-%m-%dT%H:%M:%SZ')
                nc_trg.sync()
                nc_trg.close()
                # add variables
                nc_trg = nc4.Dataset(trgFile,'a')
                qualtrg = nc_trg.createVariable('qualityIndicator','b',('time','lat','lon',),chunksizes=(1,len(y),len(x)),fill_value=128, zlib=True) # invullen
                qualtrg.setncattr('units','1')
                qualtrg.setncattr('standard_name', 'quality_probability_of_water')
                qualtrg.setncattr('long_name','Quality indicator probability of water')
                qualtrg.setncattr('grid_mapping','wgs84')
                qualtrg.setncattr('scale_factor',0.01)
                qualtrg.setncattr('add_offset',0.0)
                
                probtrg = nc_trg.createVariable('waterProbability','b',('time','lat','lon',),chunksizes=(1,len(y),len(x)),fill_value=128, zlib=True) # invullen
                probtrg.setncattr('units','1')
                probtrg.setncattr('standard_name', 'probability_of_water')
                probtrg.setncattr('long_name','Probability of water')
                probtrg.setncattr('grid_mapping','wgs84')
                probtrg.setncattr('scale_factor',0.01)
                probtrg.setncattr('add_offset',0.0)
                nc_trg.sync()
                nc_trg.close()
                
            # open the correct file and check the date
            nc_trg = nc4.Dataset(trgFile,'a')
            time = nc_trg.variables['time']
            qualtrg = nc_trg.variables['qualityIndicator']
            probtrg = nc_trg.variables['waterProbability']
            timeObjs = nc4.num2date(time[:], units=time.units, calendar=time.calendar)
            idx = np.where(np.array(timeObjs)==timeObj)[0]
            qualdata = qual[0,:,:].data
            qualdata[qual[0,:,:].mask] = qualtrg._FillValue*0.01
            probdata = prob[0,:,:].data
            probdata[prob[0,:,:].mask] = probtrg._FillValue*0.01
            qualtrg[idx,:,:] = qualdata
            probtrg[idx,:,:] = probdata
            year_old = year
            nc_trg.sync()
            nc_trg.close()

            # TODO nc_trg.time_coverage_end
            
            #pdb.set_trace()