#! /usr/bin/env python

'''
    FILE: camel_utilities.py

    PURPOSE: Provide a library of CAMEL routines to be used by ST emissivity
             applications.

    PROJECT: Land Satellites Data Systems (LSDS) Science Research and
             Development (LSRD) at the USGS EROS

    LICENSE: NASA Open Source Agreement 1.3
'''

import os
import logging
import datetime

from osgeo import gdal
import netCDF4 as nc4
import numpy as np
from scipy.interpolate import griddata

from espa import Geo
import st_utilities as util
import emissivity_utilities as emis_util


logger = logging.getLogger(__name__)

# CAMEL constants
GED_SAMPS = 1000
GED_LINES = 1000
GED_X_RES = 0.001
GED_Y_RES = 0.001
CAMEL_DEFAULT_NDVI = 0.0
CAMEL_DEFAULT_EMIS = 0.98
CAMEL_DEFAULT_EMIS_STDEV = 0.0
GED_EMIS_SOURCE = 0
CAMEL_EMIS_SOURCE = 1
WATER_EMIS_SOURCE = 2
SNOW_EMIS_SOURCE = 3
MIX_THRESHOLD = 0.75 # Skip GED if GED% (vs CAMEL) is below this; too jumbled.
NIGHT_THRESHOLD = 85 # Skip GED if solar zenith angle is above this; NDVI is
                     # unreliable here.  SZA 85 is the flight ops daylit cutoff
SKIP_GED_FILE = 'skip_ged.txt'

# CAMEL bands:
# 3.6, 4.3, 5, 5.8, 7.6, 8.3, 8.6, 9.1, 10.6, 10.8, 11.3, 12.1, 14.3
#
UM_10_6_INDEX = 8
UM_11_3_INDEX = 10
UM_12_1_INDEX = 11

camel_format = 'CAMEL.v003.{0}.{1}.0001_{2}.tif'


def get_camel_filename(espa_metadata, camel_path, emis_type):
    """ Derives the CAMEL auxiliary filename for the scene

    Args:
        espa_metadata <espa.Metadata>: XML metadata
        camel_path <str>: Path to the CAMEL archive
        emis_type <str>: Type of emissivity: mean or standard deviation

    Returns:
        <str>: Name including path for CAMEL auxiliary filename
               "CAMEL_file_not_found" if no matching file was found
    """

    # Define the CAMEL filename format
    if emis_type == 'mean':
        camel_filename_format = 'CAM5K30EM_emis_{0:4}{1:2}_V003.nc'
    else: # emis_type == 'stdev'
        camel_filename_format = \
            'CAM5K30UC_emis_uncertainty_{0:4}{1:2}_V003.nc'

    # Get the acquisition date of the scene
    acq_date = str(espa_metadata.xml_object.global_metadata.acquisition_date)
    acquisition_date = datetime.datetime.strptime(acq_date, '%Y-%m-%d')

    # Our CAMEL dataset does not cover the full range of dates we support.
    # On the most recent side of the date range, CAMEL will be updated only
    # periodically (let's say every few months or once per year).  Find the
    # closest available CAMEL month that matches the scene month.  We want
    # the season to be similar.  The earliest CAMEL month is 2000-03.  We
    # also need to deal with expected gaps in the CAMEL data.  For gaps,
    # for the same year, check month -1,+1,-2,+2, then year -1,-2 ...
    # We expect the gaps to be limited so this should be more than enough.
    # Also use the close 2 months at the start/end of the CAMEL range.

    # Check for existence of CAMEL data
    found_camel_data = False
    for year_offset in [0, -1, -2, -3, -4]:
        if found_camel_data:
            break
        for month_offset in [0, -1, 1, -2, 2]:
            camel_year = acquisition_date.year
            camel_month = acquisition_date.month + month_offset
            if acquisition_date.year < 2000: 
                # The CAMEL start month is 2000-03 so shift up to 2000
                camel_year = 2000
            if camel_year == 2000:
                if camel_month == 1 or camel_month == 2:
                    camel_year += 1
            camel_year += year_offset

            # Handle year-crossing cases
            if camel_month < 1:
                camel_year -= 1
                camel_month += 12
            if camel_month > 12:
                camel_year += 1
                camel_month -= 12

            camel_filename = camel_path + '/' + camel_filename_format.format( \
                camel_year, str(camel_month).rjust(2, '0'))
            if os.path.exists(camel_filename):
                found_camel_data = True
                return camel_filename

    return "CAMEL_file_not_found" 


def get_camel_data(camel_filename, coefficients, emis_type, band):
    """ Read and assemble global CAMEL mean or standard deviation data

    Args:
        camel_filename <str>: File with CAMEL tile data
        coefficients <CoefficientInfo>: coefficients for the math
        emis_type <str>: Type of emissivity: mean or standard deviation
        band <str> : Band to process

    Returns:
        MaskedArray : CAMEL mean or standard deviation emissivity data
        MaskedArray : grid of CAMEL latitudes
        MaskedArray : grid of CAMEL longitudes
    """

    # Open the CAMEL NetCDF file
    f = nc4.Dataset(camel_filename, 'r')

    # Read geolocation data
    lat_EMISDF = np.float32(f.variables['latitude'][:])
    lon_EMISDF = np.float32(f.variables['longitude'][:])

    # Read fill value
    if emis_type == 'mean':
        cemis_fill = f.variables['camel_emis']._FillValue

        if band == 'band11':
            # Read CAMEL emissivity for 12.1 um band for all latitude/longitudes
            cemis1 = np.float32(f.variables['camel_emis'][:,:,UM_12_1_INDEX])
            cemis1[cemis1 == cemis_fill] = np.NaN
            cemis1[cemis1 < 0] = np.NaN

            # Compute equivalent emissivity for Landsat thermal band
            camel_emis = coefficients.camel_1 * cemis1 + coefficients.camel_2
        else:
            # Read CAMEL emissivity for 10.6 um band for all latitude/longitudes
            cemis1 = np.float32(f.variables['camel_emis'][:,:,UM_10_6_INDEX])
            cemis1[cemis1 == cemis_fill] = np.NaN
            cemis1[cemis1 < 0] = np.NaN

            # Read CAMEL emissivity for 11.3 um band for all latitude/longitudes
            cemis3 = np.float32(f.variables['camel_emis'][:,:,UM_11_3_INDEX])
            cemis3[cemis3 == cemis_fill] = np.NaN
            cemis3[cemis3 < 0] = np.NaN

            # Compute equivalent emissivity for Landsat thermal band
            camel_emis = coefficients.camel_1 * cemis1 \
                       + coefficients.camel_2 * cemis3 + coefficients.camel_3
    else: # emis_type == 'stdev'
        cemis_fill = f.variables['total_uncertainty']._FillValue

        if band == 'band11':
            # Read CAMEL emissivity stdev for 12.1 um band for all latitude/
            # longitudes
            camel_emis = np.float32(f.variables['total_uncertainty']
                                               [:,:,UM_12_1_INDEX])
        else:
            # Read CAMEL emissivity stdev for 10.6 um band for all latitude/
            # longitudes
            camel_emis = np.float32(f.variables['total_uncertainty']
                                               [:,:,UM_10_6_INDEX])
        camel_emis[camel_emis == cemis_fill] = np.NaN
        camel_emis[camel_emis < 0] = np.NaN

    # Make grid of latitude and longitude values
    [latg_EMISDF, long_EMISDF] = np.meshgrid(lat_EMISDF, lon_EMISDF,
                                             indexing='ij')

    return camel_emis, latg_EMISDF, long_EMISDF


def make_camel_file(camel_emis, latg_EMISDF, long_EMISDF, max_lat, max_lon,
                    min_lat, min_lon, wkt, emis_file_name, no_data_value,
                    first_band, emis_type):
    """Generate a tile for emissivity mean or standard deviation from CAMEL
       data

    Args:
        camel_emis <MaskedArray> : CAMEL mean or standard deviation emissivity
                                   data
        latg_EMISDF <MaskedArray> : grid of CAMEL latitudes
        long_EMISDF <MaskedArray> : grid of CAMEL longitudes
        max_lat <int> : Maximum latitude of corresponding GED tiles' corner
        max_lon <int> : Maximum longitude of corresponding GED tiles' corner
        min_lat <int> : Minimum latitude of corresponding GED tiles' corner
        min_lon <int> : Minimum longitude of corresponding GED tiles' corner
        wkt <str>: Well-Known-Text describing the projection
        emis_file_name <str> : Name of emissivity file to make
        no_data_value <float>: Value to use for fill
        first_band <bool>: Is this the first band to be processed?
        emis_type <str>: Type of emissivity: mean or standard deviation
    """

    # Extend in both directions to cover last GED tiles, not just the corner
    maxlatA = max_lat
    maxlonA = max_lon + 1
    minlatA = min_lat - 1
    minlonA = min_lon

    # The number of ASTER GED sized tiles in both dimensions to cover the scene
    lat_difference = abs(maxlatA - minlatA)
    lon_difference = abs(maxlonA - minlonA)

    # Fill in latitudes/longitudes between the corners
    num_lines = lat_difference * GED_LINES
    num_samples = lon_difference * GED_SAMPS
    latsA = np.linspace(minlatA, maxlatA, num_lines)
    lonsA = np.linspace(minlonA, maxlonA, num_samples)

    [lat_g, lon_g] = np.meshgrid(latsA, lonsA, indexing='ij')
    lat_g = np.flip(lat_g, 0)

    # Crop to ASTER grid.  Find indexes to remove
    I1 = np.where((latg_EMISDF[:, 0] < minlatA - 1) \
                | (latg_EMISDF[:, 0] > maxlatA + 1))

    # CAMEL longitudes are -179.975 to 179.975.  The griddata interpolation
    # only works within the boundary of the points, so for tiles along the
    # antimeridian, we need to add points from the other side.  But add a 
    # full degree since it improves the interpolation
    I2 = np.where((long_EMISDF[0, :] < minlonA - 1) \
                | (long_EMISDF[0, :] > maxlonA + 1))

    # Initialize cropped variables
    latg_cut = latg_EMISDF;
    long_cut = long_EMISDF;
    camel_emis_cut = np.ma.filled(camel_emis, np.nan)

    # Remove the unneeded indexes along both axes
    latg_cut = np.delete(latg_cut, I1, 0)
    latg_cut = np.delete(latg_cut, I2, 1)
    long_cut = np.delete(long_cut, I1, 0)
    long_cut = np.delete(long_cut, I2, 1)
    camel_emis_cut = np.delete(camel_emis_cut, I1, 0)
    camel_emis_cut = np.delete(camel_emis_cut, I2, 1)

    # Check for missing or bad values and set to default
    fill_locations = np.where(np.isnan(camel_emis_cut))
    if emis_type == 'mean':
        camel_emis_cut[fill_locations] = CAMEL_DEFAULT_EMIS
    else: #emis_type == 'stdev'
        camel_emis_cut[fill_locations] = CAMEL_DEFAULT_EMIS_STDEV

    # Interpolate CAMEL points to fill ASTER grid
    latg_cut_flat = latg_cut.flatten()
    long_cut_flat = long_cut.flatten()
    camel_emis_cut_flat = camel_emis_cut.flatten()
    file_emis = griddata((latg_cut_flat, long_cut_flat), camel_emis_cut_flat,
                         (lat_g, lon_g), 'cubic')

    # Define transformation needed to generate the CAMEL file
    transform = [min_lon, GED_X_RES, 0, max_lat, 0, -GED_Y_RES]

    # Convert float64 memory to float32 before writing to float32 file
    file_emis = np.float32(file_emis)

    # Create the estimated EMIS raster output file matching extent of ASTER GED
    # tiles
    if emis_type == 'mean':
        logger.debug('Creating a CAMEL estimated EMIS file {}'.
                     format(emis_file_name))
    else: #emis_type == 'stdev'
        logger.debug('Creating a CAMEL estimated EMIS standard deviation file {}'.
                     format(emis_file_name))

    # This reverses order of nodata and WriteArray steps compared to
    # generate_raster_file preventing 0 values from becoming nodata
    driver = gdal.GetDriverByName('GTiff')
    raster = driver.Create(emis_file_name, num_samples, num_lines, 1, gdal.GDT_Float32)

    raster.SetGeoTransform(transform)
    raster.SetProjection(wkt)
    raster.GetRasterBand(1).SetNoDataValue(no_data_value)
    raster.GetRasterBand(1).WriteArray(file_emis)
    raster.FlushCache()

    # Cleanup memory
    del raster

    del file_emis


def generate_camel_file(camel_filename, camel_tile_list, coefficients, wkt,
                        no_data_value, first_band, emis_type, band):
    """Generate tiles covering ASTER GED extent for emissivity mean or
       standard deviation from CAMEL data

    Args:
        camel_filename <str>: File with CAMEL tile data
        camel_tile_list <[str]> : List of CAMEL tiles
        coefficients <CoefficientInfo>: coefficients for the math
        wkt <str>: Well-Known-Text describing the projection
        no_data_value <float>: Value to use for fill
        first_band <bool>: Is this the first band to be processed?
        emis_type <str>: Type of emissivity: mean or standard deviation
        band <str> : Band to process

    Returns:
        <str>: Emissivity mean or standard deviation filename
    """

    # If there are tiles that require CAMEL data, do initial processing to set
    # up CAMEL data
    if camel_tile_list:
        (camel_emis, latg_EMISDF, long_EMISDF) = \
            get_camel_data(camel_filename, coefficients, emis_type, band)

    # Each tile is 1x1 degree, but we will generate them as a single band.
    # Find the range of latitudes and longitudes for this
    lat_list = []
    lon_list = []
    for camel_tile_loc in camel_tile_list:
        (lat, lon) = camel_tile_loc.split('.')
        lat_list.append(int(lat))
        lon_list.append(int(lon))

    max_lat = max(lat_list)
    max_lon = max(lon_list)
    min_lat = min(lat_list)
    min_lon = min(lon_list)

    # Handle the case of antimeridian crossing or border (since we need to
    # cross the antimeridian to interpolate the border ones)
    if max_lon == 179 or min_lon == -180:
        # Make a single contiguous range for the longitudes
        lon_list = [x + 360 if x <= 0 else x for x in lon_list]

        # Find the min and max of the longitudes of the contiguous range
        max_lon = max(lon_list)
        min_lon = min(lon_list)

        # Adjust the CAMEL longitudes accordingly
        long_EMISDF[long_EMISDF <= 0] += 360

    # Build the output file name
    if emis_type == 'mean':
        emis_file_name = camel_format.format(max_lat, min_lon, 'emis')
    else: # emis_type == 'stdev'
        emis_file_name = camel_format.format(max_lat, min_lon, 'emis_stdev')

    # Create emissivity file using CAMEL data
    make_camel_file(camel_emis, latg_EMISDF, long_EMISDF, max_lat, max_lon,
                    min_lat, min_lon, wkt, emis_file_name, no_data_value,
                    first_band, emis_type)

    return emis_file_name