#! /usr/bin/env python
# encoding: utf-8
'''
RichF - check that you can see this line

@author: nick.bearson@ssec.wisc.edu

Copyright (c) 2012 University of Wisconsin SSEC.  All rights reserved.
'''

import os
from optparse import OptionParser
import re
from datetime import datetime, timedelta
import glob
import logging

import numpy as np
from numpy import ma

from pypre.variables.Variable import Variable

import pypre.patterns.calc_patterns as cp
import pypre.patterns.ancil_patterns as ap
import pypre.patterns.viirs_patterns as vp

from pypre.files.HDF4File import HDF4File
from pypre.files.HDF5File import HDF5File
from pypre.files.GribFile import GribFile

from pypre.util import make_filename, get_file_date

import pyvcm.vcm_config as vcmc
from   pyvcm.vcm_include import imgr_bands, FLOAT_FILL

LOG = logging.getLogger(__name__)

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

def vcm_gatherer(datetime, outputFile):

  #################################################
  # GEOLOCATION DATA:
  viirs_re = vcmc.get('inputs', 'VIIRS')
  geolocFilename = glob.glob(make_filename(viirs_re, datetime, sds_tag="GMTCO"))[0]
  LOG.info("Gathering: %s" % geolocFilename)
  geolocFile = HDF5File(geolocFilename)

  lat = vp.viirs_latitude(geolocFile, outputFile)
  lon = vp.viirs_longitude(geolocFile, outputFile)

  import numpy.ma as ma

  latMask = (lat.data < -180) | (lat.data > 180)
  lonMask = (lon.data < -360) | (lon.data > 360)
  lat.data = ma.masked_array(lat.data, mask=latMask)
  lon.data = ma.masked_array(lon.data, mask=lonMask)

  
  solz = vp.viirs_solar_zenith(geolocFile, outputFile)
  senz = vp.viirs_sensor_zenith(geolocFile, outputFile)
  sola = vp.viirs_solar_azimuth(geolocFile, outputFile)
  sena = vp.viirs_sensor_azimuth(geolocFile, outputFile)

  nscans = vp.viirs_nscans(geolocFile, outputFile)
  height = vp.viirs_height(geolocFile, outputFile)

  M_shape = lat.data.shape
  I_shape = tuple([2 * x for x in M_shape])

  #################################################
  # VIIRS IMGR:
  for units in imgr_bands.keys():
    for res in imgr_bands[units].keys():
      for sds in imgr_bands[units][res].keys():
#        try:
        imgr_fname = glob.glob(make_filename(viirs_re, datetime, sds_tag=sds))[0]
        LOG.info("Gathering: %s" % imgr_fname) 
        imgrFile = HDF5File(imgr_fname)
        data = vp.viirs_imgr_data(imgrFile, units, outputFile)
        imgr_bands[units][res][sds] = data


###################################################
#
# Now that all the bands are online, there's no good reason for a missing input band... is there?
#        except IOError as detail:
#          print "IO error (file likely doesn't exist) - setting", units, " for ", sds, " to None"
#          if res is "M":
#            blank = np.ones(M_shape) * FLOAT_FILL
#            imgr_bands[units][res][sds] = Variable(data=blank)
#            outputFile.writeVariable(imgr_bands[units][res][sds], writeVarName=(sds + "_" + units))
#          elif res is "I":
#            blank = np.ones(I_shape) * FLOAT_FILL
#            imgr_bands[units][res][sds] = Variable(data=blank)
#            outputFile.writeVariable(imgr_bands[units][res][sds], writeVarName=(sds + "_" + units))
#          else:
#            assert(False)
#
#
###################################################
  # ANCIL:

  igbpFilename = make_filename(vcmc.get('inputs', 'IGBP'), datetime)
  LOG.info("Gathering: %s" % igbpFilename) 
  igbpFile = HDF4File(igbpFilename)
  igbp = ap.igbp(igbpFile, lon, lat)
  outputFile.writeVariable(igbp, writeVarName="MODIS IGBP")

  igbp.data = cp.igbp_modis2viirs(igbp.data)
  outputFile.writeVariable(igbp, writeVarName="VIIRS IGBP")

  # TODO: FIND REPLACEMENT FOR MAKE_FILENAME_NDVI AND MAKE_FILENAME_GDAS
  ndvi_date = get_file_date(datetime, "ndvi")
  ndviFilename = make_filename(vcmc.get('inputs', 'NDVI'), ndvi_date)
  LOG.info("Gathering: %s" % ndviFilename)
  ndviFile = HDF4File(ndviFilename)
  ndvi = ap.ndvi(ndviFile, lat, lon, outputFile)

  gdas_early, gdas_late = get_file_date(datetime, "gdas")

  gdasEFilename = make_filename(vcmc.get('inputs', 'GDAS'), gdas_early)
  gdasLFilename = make_filename(vcmc.get('inputs', 'GDAS'), gdas_late)

  LOG.info("Gathering: %s" % gdasEFilename)
  gdasEFile = GribFile(gdasEFilename)

  e_ph2o = ap.precipitable_water(gdasEFile, 0, lon, lat, None, 1)
  e_wind = ap.wind_speed(gdasEFile, lon, lat, None, 1)
  e_sfct = ap.surface_temp(gdasEFile, lon, lat, None, 1)

  LOG.info("Gathering: %s" % gdasLFilename)
  gdasLFile = GribFile(gdasLFilename)

  l_ph2o = ap.precipitable_water(gdasLFile, 0, lon, lat, None)
  l_wind = ap.wind_speed(gdasLFile, lon, lat, None)
  l_sfct = ap.surface_temp(gdasLFile, lon, lat, None, 1)

  # just adding this to the output file for cross-checking, not actually used
#  lsm = ap.land_sea_mask(gdasEFile, lon, lat, outputFile)

  # masked sfct
#  msfct = ap.masked_sfct(gdasEFile, lon.data, lat.data, outputFile, 1)


  LOG.info("Temporally interpolating GDAS fields.")
  ph2o   = t_interp(e_ph2o.data, gdas_early, l_ph2o.data, gdas_late, datetime)
  wind   = t_interp(e_wind.data, gdas_early, l_wind.data, gdas_late, datetime)
#  sfct   = t_interp(e_sfct.data, gdas_early, l_sfct.data, gdas_late, datetime)

  outputFile.writeData("PrecipitableWater_L0", ph2o)
  outputFile.writeData("10 metre wind speed", wind)
#  outputFile.writeData("SurfaceTemperature", sfct)

  # do DEM_LW
  demlwFilename = vcmc.get('inputs', 'DEM_LW')
  LOG.info("Gathering: %s" % demlwFilename)
  demlwFile = HDF5File(demlwFilename)
  dem_lw = ap.dem_land_water(demlwFile, lon, lat, outputFile)

  niseFilename = make_filename(vcmc.get('inputs', 'NISE'), datetime)
  LOG.info("Gathering: %s" % niseFilename)
  niseFile = HDF4File(niseFilename)
  nise = ap.nise(niseFile, lon, lat, outputFile)

  #################################################
  # CALC:
#  _, e_sfct = cp.sfct_interpd(gdasEFile, igbpFile, lat.data, lon.data)
#  _, l_sfct = cp.sfct_interpd(gdasLFile, igbpFile, lat.data, lon.data)
  sfct   = t_interp(e_sfct.data, gdas_early, l_sfct.data, gdas_late, datetime)
  outputFile.writeData("SurfaceTemperature", sfct)
#  outputFile.writeData("SurfaceTemperature_Split", sfct2)

  qstLwm = cp.qstLwm(igbp, dem_lw, nscans, outputFile)

  snowice = calculate_snow_ice(nise, qstLwm)
  snowice = Variable(data=snowice, varName="snowice")
  outputFile.writeVariable(snowice)

CM_SNOW=1
def calculate_snow_ice(nise, qstLwm):
  nise = nise.data
  qstlwm = qstLwm.data

  snowice = np.zeros(nise.shape, dtype='uint8')

  idx101 = (nise == 101)
  idx103 = (nise == 103)
  idx104 = (nise == 104)

  idxLandIceGT25 = ( (nise > 25) & (nise < 101) & ( (qstlwm < 17)  | (qstlwm == 19) ) )
  idxWaterIceGT25 = ( (nise > 25) & (nise < 102) & ( (qstlwm == 17) | (qstlwm == 18) ) )

  idxAll = idx101 | idx103 | idx104 | idxLandIceGT25 | idxWaterIceGT25
  snowice[idxAll] = 1

  return snowice



"""
Temporal interpolation

Linear for the moment!
"""
import time as _time
def t_interp(grid0, date0, grid1, date1, key_date):
#  print "\n+ t_interp:"
  time0    = _time.mktime(date0.timetuple())
  time1    = _time.mktime(date1.timetuple())
  key_time = _time.mktime(key_date.timetuple())
  res = grid0 + (key_time - time0) * (grid1 - grid0) / (time1 - time0)
#  print grid0
#  print grid1
#  print res
  return res

"""
  if(
   //Permanent Ice
  nise_val == 101 || 
   //Dry Snow
  nise_val == 103 || 
   //Wet Snow
  nise_val == 104 || 
   //Transient sea ice > 25% concentration, AND (Land OR Coastal)
  (nise_val > 25 && nise_val < 101 && 
          ( *(vcmDataType->qstLwm->data[0] + idx) < 17 || *(vcmDataType->qstLwm->data[0] + idx) == 19 )) || 
   //Transient sea ice > 25% concentration, AND Permanent Ice AND(Sea OR Inland Water)
  (nise_val > 25 && nise_val < 102 && 
          ( *(vcmDataType->qstLwm->data[0] + idx) == 17 || *(vcmDataType->qstLwm->data[0] + idx) == 18 )) )
{
  *(vcmDataType->snow->data[0] + idx) = CM_SNOW;
}
"""

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

def _check_args(args):
  if len(args) != 2:
    print "Invalid arguments."
    print _usage()
    exit(-1)
#  if not os.access(args[0], os.F_OK):
#    print "Input file doesn't exist."
#    print _usage()
#    exit(-1)

def _usage():
  return """Usage: vcm_gatherer geolocation.h5 output.h5"""

def _handle_args():
  parser = OptionParser(usage=_usage())
#  parser.add_option("-g", "--geofile", action="store", type="string", dest="geoFile", help="geolocation file to use")
#  parser.add_option("-a", "--ancFile", action="store", type="string", dest="ancFile", help="ancilliary  h5 file")
#  parser.add_option("-c", "--cmoFile", action="store", type="string", dest="cmoFile", help="cloudmask output h5 file")
  return parser.parse_args()

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

def main():
  logging.basicConfig(level=logging.INFO)

  # Handle arguments
  (options, args) = _handle_args()

  _check_args(args)

  f_geolocation = args[0]
  f_output = args[1]

  geolocFilename = os.path.abspath(f_geolocation)
  geolocBasename, geolocExtension = os.path.splitext(geolocFilename)

  match = re.search('.*d(\d{4})(\d{2})(\d{2})_t(\d{2})(\d{2})(\d{2}).*', geolocFilename)
  if not match:
    raise RuntimeError, "Can't extract time info from geolocation filename."

  year   = int(match.group(1))
  month  = int(match.group(2)) 
  day    = int(match.group(3))
  hour   = int(match.group(4))
  minute = int(match.group(5))
  second = int(match.group(6))

  wantedDatetime = datetime(year=year, month=month, day=day, hour=hour, minute=minute, second=second)

  print "GeoLocation File: " + geolocFilename

  outFilename = f_output # TODO: make variable names more consistent
  if os.access(outFilename, os.F_OK):
    os.remove(outFilename)

  # Create file object for output file:
  print "Out: " + outFilename

  outputFile = HDF5File(outFilename)


  vcm_gatherer(wantedDatetime, outputFile)


if __name__ == "__main__":
  main()