#!/usr/bin/python

#

# Early version of assorted Python code to download MODIS 11A1 (daily)

# data associated with specified dates and tiles, then interface with

# GRASS to calculate temporally averaged LST in a way that accounts for

# QC flags.

#

# TODO:

#  - functionalize to encapsulate high level procedural steps

#  - create a 'main' function to allow script to be run as a command

#    that can accept arguments?

#  - put downloaded data somewhere other than working directory?

#  - write code to set up and take down a temporary GRASS location/mapset

#  - write code to export climatology rasters to file (geotiff? compressed?)

#  - calculate other aggregate stats? (stddev, ...)

#  - deal with nightly LST too?

#  - deal with more than just first two bits of QC flags?

#     e.g.: r.mapcalc 'qc_level2 = (qc>>2) & 0x03'

#

# Jim Regetz

# NCEAS

# Created on 16-May-2012

import os

import datetime

import ftplib

import grass.script as gs

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

# helper functions

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

# quick function to reformat e.g. 200065 to 2000.03.05

def yj_to_ymd(year, doy):

    date = datetime.datetime.strptime('%d%d' % (year, doy), '%Y%j')

    return date.strftime('%Y.%m.%d')

# quick function to return list of dirs in wd

def list_contents():

    listing = []

    ftp.dir(listing.append)

    contents = [item.split()[-1] for item in listing[1:]]

    return contents

# create raster of pixelwise N for a list of input rasters

def calc_nobs(maplist, name, overwrite=False):

    denominator = '(%s)' % '+'.join(['if(!isnull(%s))' % m

        for m in maplist])

    gs.mapcalc('nobs_%s = %s' % (name, denominator), overwrite=overwrite)

# create raster of pixelwise means for a list of input rasters

def calc_mean(maplist, name, overwrite=False):

    numerator = '(%s)' % '+'.join(['if(isnull(%s), 0, %s)' % (m, m)

        for m in maplist])

    denominator = '(%s)' % '+'.join(['if(!isnull(%s))' % m

        for m in maplist])

    gs.mapcalc('mean_%s = %s/%s' % (name, numerator, denominator),

        overwrite=overwrite)

# ...combined version of the above two functions...

def calc_clim(maplist, name, overwrite=False):

    denominator = '(%s)' % '+'.join(['if(!isnull(%s))' % m

        for m in maplist])

    gs.mapcalc('nobs_%s = %s' % (name, denominator), overwrite=overwrite)

    numerator = '(%s)' % '+'.join(['if(isnull(%s), 0, %s)' % (m, m)

        for m in maplist])

    #gs.mapcalc('mean_%s = round(float(%s)/nobs_%s)' % (name, numerator, name),

    gs.mapcalc('mean_%s = %s/nobs_%s' % (name, numerator, name),

        overwrite=overwrite)

def load_qc_adjusted_lst(hdf):

    """Load LST_Day_1km into GRASS from the specified hdf file, and

    nullify any cells for which QA flags indicate anything other than

    high quality.

    Argument:

    hdf -- local path to the 11A1 HDF file

    Returns the name of the QC-adjusted LST map created in GRASS.

    """

    lstname =  'LST_' + '_'.join(os.path.basename(hdf).split('.')[1:3])

    # read in lst grid

    gs.run_command('r.in.gdal',

        input='HDF4_EOS:EOS_GRID:%s:MODIS_Grid_Daily_1km_LST:LST_Day_1km' % hdf,

        output=lstname)

    # read in qc grid

    gs.run_command('r.in.gdal',

        input = 'HDF4_EOS:EOS_GRID:%s:MODIS_Grid_Daily_1km_LST:QC_Day' % hdf,

        output = 'qc_this_day')

    gs.run_command('g.region', rast=lstname)

    # null out any LST cells for which QC is not high [00]

    gs.mapcalc('${lst} = if((qc_this_day & 0x03)==0, ${lst}, null())',

        lst=lstname, overwrite=True)

    # clean up

    gs.run_command('g.remove', rast='qc_this_day')

    # return name of qc-adjusted LST raster in GRASS

    return lstname

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

# set up a (temporary?) GRASS data store

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

## TODO ##

# hack to fix datum???

gs.run_command('g.proj', flags='c',

    proj4='+proj=sinu +a=6371007.181 +b=6371007.181 +ellps=sphere')

#   proj4='+proj=sinu +R=6371007.181 +nadgrids=@null +wktext')

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

# test download & process

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

tile = 'h09v04'

year = 2000

start_doy = 122

end_doy = 152

# connect to data pool and change to the right directory

ftp = ftplib.FTP('e4ftl01.cr.usgs.gov')

ftp.login('anonymous', '')

ftp.cwd('MOLT/MOD11A1.005')

# make list of daily directories to traverse

day_dirs = listdirs()

days_to_get = [yj_to_ymd(year, x) for x in range(start_doy, end_doy+1)]

if not all [day in day_dirs for day in days_to_get]:

    error

# using v004 data, this was taking 60 seconds for the 8 hdf/xml pairs on

# atlas [16-May-2012] ... each individual hdf is ~23MB

hdfs = list()

for day in days_to_get:

    ftp.cwd(day)

    files_to_get = [file for file in list_contents()

        if tile in file and file[-3:]=='hdf']

    if len(files_to_get)==0:

        # no file found -- print message and move on

        print 'no hdf found on server for tile', tile, 'on', day

        continue

    elif 1<len(files_to_get):

        # multiple files found! -- just use the first...

        print 'multiple hdfs found on server for tile', tile, 'on', day

    file = files_to_get[0]

    ftp.retrbinary('RETR %s' % file, open(file, 'wb').write)

    hdfs.append(file)

    ftp.cwd('..')

# politely disconnect

ftp.quit()

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

gs.os.environ['GRASS_OVERWRITE'] = '1'

# load and qc-adjust all daily LSTs, return list of the GRASS map names

LST = [load_qc_adjusted_lst(hdf) for hdf in hdfs]

# calculate mean and nobs rasters

calc_clim(LST, 'LST_%s_%d_%03d_%03d' % (tile, year, start_doy, end_doy))

# clean up

gs.run_command('g.remove', rast=','.join(LST))

gs.os.environ['GRASS_OVERWRITE'] = '0'

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

#

# test 31-day mean with local files

#

# Note that some statements below are redundant with (and possibly

# slightly different) from code above -- still need to consolidate

start_doy = 122

end_doy = 152

gs.os.environ['GRASS_OVERWRITE'] = '1'

ordir = '/home/layers/data/climate/MOD11A1.004-OR-orig-2000'

LST = list()

for doy in range(start_doy, end_doy+1):

    fileglob = 'MOD11A1.A%d%03d.%s*hdf' % (year, doy, tile)

    try:

        hdfpath = glob.glob(os.path.join(ordir, fileglob))[0]

    except:

        print '*** problem finding', fileglob, '***'

        continue

    hdfname = os.path.basename(hdfpath)

    LST.append('LST_' + '_'.join(hdfname.split('.')[1:3]))

    # TODO: handle failures here...

    gs.run_command('r.in.gdal',

        input='HDF4_EOS:EOS_GRID:%s:MODIS_Grid_Daily_1km_LST:LST_Day_1km' % hdfpath,

        output=LST[-1])

gs.os.environ['GRASS_OVERWRITE'] = '0'

gs.run_command('g.region', rast=LST[0])

#numerator = '(%s)' % '+'.join(['if(isnull(%s), 0, %s)' % (lst, lst) for lst in LST])

#denominator = '(%s)' % '+'.join(['if(!isnull(%s))' % lst for lst in LST])

#gs.mapcalc('LST_mean_%s = %s / %s' % (tile, numerator, denominator))

calc_mean(LST, 'LST_mean_%s_%d_%03d_%03d' % (tile, year, start_doy, end_doy))

# clean up

gs.run_command('g.remove', rast=','.join(LST))