#! /usr/bin/env python

'''
    File: st_split_window.py

    Purpose: Run the core processing of the ST Split Window algorithm.

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

    License: NASA Open Source Agreement 1.3

'''

import os
import sys
import logging
import datetime
import json
from argparse import ArgumentParser
from collections import namedtuple
import st_utilities as util
from st_split_window_tpw_pred import main as pred_tpw

import numpy as np
import astropy.convolution.convolve as ap_convolve
from astropy.convolution import Box2DKernel

from lxml import objectify
from osgeo import gdal, osr

from espa import Metadata, Sys, Geo, Dataset
from st_exceptions import MissingBandError, InvalidCoefficientsError
import st_split_window_defines as sw_def
import camel_utilities as cm_util

ThermalConstantInfo = namedtuple('ThermalInfo', ('k1', 'k2'))

# Global variables
logger = logging.getLogger(__name__)


def retrieve_command_line_arguments():
    """Read arguments from the command line

    Returns:
        args <arguments>: The arguments read from the command line
    """

    parser = ArgumentParser(description='Creates the Surface Temperature' +
                                        ' band using Split Window Algorithm')

    parser.add_argument('--version',
                        action='version',
                        version=util.Version.split_window_version_text())

    parser.add_argument('--xml',
                        action='store', dest='xml_filename',
                        required=True, default=None,
                        help='The XML metadata file to use')

    parser.add_argument('--data_path',
                        action='store', dest='data_path',
                        required=True, default=None,
                        help='Specify the ST Data directory')

    parser.add_argument('--scale',
                        action='store', dest='scale', type=float,
                        required=False, default=util.DEFAULT_SCALE,
                        help='The scale factor for ST products')

    parser.add_argument('--offset',
                        action='store', dest='offset', type=float,
                        required=False, default=util.DEFAULT_OFFSET,
                        help='The offset value for ST products')

    args = parser.parse_args()

    return args


def get_band_filename(espa_metadata, band_name):
    """Reads required information from the metadata XML file

    Args:
        espa_metadata <espa.Metadata>: XML metadata
        band_name <str>: Name of band to extract metadata information from

    Returns:
        filename <str>: Filename for the requested band
    """

    filename = None

    # Find the band to extract information from
    for band in espa_metadata.xml_object.bands.band:
        if band.get('name') == band_name:
            filename = str(band.file_name)
            break

    # Error if we didn't find the required band in the data
    if filename is None:
        raise MissingBandError('Failed to find band ' + band_name +
                               ' in the input data')

    return filename


def retrieve_scaling_parms(espa_metadata, band_type, band_name):
    """Reads required information from the metadata XML file

    Args:
        espa_metadata <espa.Metadata>: XML metadata
        band_type <str>: "BT" or "THERMAL" band
        band_name <str>: Name of band to extract metadata information from

    Returns:
        Scale and offset values
    """

    scale = None
    offset = None

    # Find the band to extract information from
    for band in espa_metadata.xml_object.bands.band:
        if band.get('name') == band_name:
            if band_type == "BT":
                scale = band.get('scale_factor')
                offset = band.get('add_offset')
                break
            else: # "THERMAL"
                scale = band.radiance.get('gain')
                offset = band.radiance.get('bias')
                break

    # Error if we didn't find the required band or the band's information
    if scale is None or offset is None:
        raise MissingBandError('Failed to find the band in the input data '
                               ' or values for the band')

    return (float(scale), float(offset))


def read_metadata(espa_metadata):
    """Reads required information from the metadata XML file

    Args:
        espa_metadata <espa.Metadata>: XML metadata

    Returns:
        emis_b10_band_filename <str>: emissivity band 10 filename
        emis_b11_band_filename <str>: emissivity band 11 filename
        emis_stdev_b10_band_filename <str>: emissivity stdev band 10 filename
        emis_stdev_b11_band_filename <str>: emissivity stdev band 11 filename
        emis_source_band_filename <str>: emissivity source band filename
        qa_pixel_band_filename <str>: QA band filename
        thermal_b10_band_filename <str>: thermal band 10 filename
        thermal_b11_band_filename <str>: thermal band 11 filename
        thermal_b10_scale <float>: scale factor for thermal band 10
        thermal_b10_offset <float>: offset for thermal band 11
        thermal_b11_scale <float>: scale factor for thermal band 10
        thermal_b11_offset <float>: offset for thermal band 11
    """

    emis_b10_band_filename = get_band_filename(espa_metadata,
                                               sw_def.EMIS_BAND10_NAME)
    emis_b11_band_filename = get_band_filename(espa_metadata,
                                               sw_def.EMIS_BAND11_NAME)
    emis_stdev_b10_band_filename = get_band_filename(espa_metadata,
                                               sw_def.EMIS_STDEV_BAND10_NAME)
    emis_stdev_b11_band_filename = get_band_filename(espa_metadata,
                                               sw_def.EMIS_STDEV_BAND11_NAME)
    emis_source_band_filename = get_band_filename(espa_metadata,
                                               sw_def.EMIS_SOURCE_NAME)
    qa_pixel_band_filename = get_band_filename(espa_metadata,
                                               sw_def.QA_PIXEL_NAME)
    thermal_b10_band_filename = get_band_filename(espa_metadata,
                                                  sw_def.BAND10_NAME)
    (thermal_b10_scale, thermal_b10_offset) = retrieve_scaling_parms(
                                             espa_metadata, "THERMAL",
                                             sw_def.BAND10_NAME)
    thermal_b11_band_filename = get_band_filename(espa_metadata,
                                                  sw_def.BAND11_NAME)
    (thermal_b11_scale, thermal_b11_offset) = retrieve_scaling_parms(
                                             espa_metadata, "THERMAL",
                                             sw_def.BAND11_NAME)

    return(emis_b10_band_filename, emis_b11_band_filename,
           emis_stdev_b10_band_filename, emis_stdev_b11_band_filename,
           emis_source_band_filename, qa_pixel_band_filename,
           thermal_b10_band_filename, thermal_b11_band_filename,
           thermal_b10_scale, thermal_b10_offset, thermal_b11_scale,
           thermal_b11_offset)


def write_product(samps, lines, transform, wkt, no_data_value, data_type,
                  filename, image_data):
    """Creates a product file

    Args:
        samps <int>: Samples in the data
        lines <int>: Lines in the data
        transform <2x3:float>: GDAL Affine transformation matrix
                               [0] - Map X of upper left corner
                               [1] - Pixel size in X direction
                               [2] - Y rotation
                               [3] - Map Y of upper left corner
                               [4] - X rotation
                               [5] - Pixel size in Y direction
        wkt <str>: Well-Known-Text describing the projection
        no_data_value <float>: Value to use for fill
        data_type <str>: Data type field in metadata
        filename <str>: Full path for the output file to create
        image_data <numpy.2darray>: Binary image data to be output
    """

    logger.info('Creating {0}'.format(filename))

    Geo.generate_raster_file(gdal.GetDriverByName('ENVI'),
                             filename,
                             image_data,
                             samps,
                             lines,
                             transform,
                             wkt,
                             no_data_value,
                             data_type) # for now in dev ST UNC is float here

    hdr_filename = Sys.replace_extension(filename, '.hdr')
    logger.info('Updating {0}'.format(hdr_filename))
    util.ST_Geo.update_envi_header(hdr_filename, no_data_value)

    # Remove the *.aux.xml file generated by GDAL
    aux_filename = Sys.replace_extension(filename, '.img.aux.xml')
    if os.path.exists(aux_filename):
        os.unlink(aux_filename)


def add_band_to_xml(espa_metadata, filename, product, name, category,
                    data_type, short_name, long_name, data_units,
                    valid_range_min, valid_range_max, no_data_value,
                    scale, offset):
    """Adds a band to the Metadata XML file

    Args:
        espa_metadata <espa.Metadata>: XML metadata information
        filename <str>: Full path for the output file to create
        product <str>: Product field in metadata
        name <str>: Name field in metadata
        category <str>: Category field in metadata
        data_type <str>: Data type field in metadata
        short_name <str>: Short name field in metadata
        long_name <str>: Long name field in metadata
        data_units <str>: Data units field in metadata
        valid_range_min <str>: Valid range min field in metadata
        valid_range_max <str>: Valid range max field in metadata
        no_data_value <float>: Value to use for fill
        scale <float>: scale_factor for ST products
        offset <float>: add_offset for ST products
    """

    logger.info('Adding {0} to Metadata XML'.format(filename))

    # Create an element maker
    maker = objectify.ElementMaker(annotate=False, namespace=None, nsmap=None)

    source_product = 'level_1_thermal'

    # Find the level 1 thermal band to use for the specific band details
    base_band = util.get_thermal_band_metadata(espa_metadata)

    st_band = maker.band()
    st_band.set('product', product)
    st_band.set('source', source_product)
    st_band.set('name', name)
    st_band.set('category', category)
    st_band.set('data_type', data_type)
    st_band.set('nlines', base_band.attrib['nlines'])
    st_band.set('nsamps', base_band.attrib['nsamps'])
    st_band.set('fill_value', str(no_data_value))
    st_band.set('scale_factor', str(scale))

    if product == sw_def.ST_PRODUCT:
        st_band.set('add_offset', str(offset))

    st_band.short_name = maker.element(short_name)
    st_band.long_name = maker.element(long_name)
    st_band.file_name = maker.element(filename)

    st_band.pixel_size = base_band.pixel_size

    st_band.resample_method = maker.element('none')
    st_band.data_units = maker.element(data_units)

    st_band.valid_range = maker.element()
    st_band.valid_range.set('min', valid_range_min)
    st_band.valid_range.set('max', valid_range_max)

    st_band.app_version = maker.element(
        util.Version.split_window_app_version())

    # Get the production date and time in string format
    # Strip the microseconds and add a Z
    date_now = ('{0}Z'.format(datetime.datetime.utcnow()
                              .strftime('%Y-%m-%dT%H:%M:%S')))
    st_band.production_date = maker.element(date_now)

    # Append the band to the XML
    espa_metadata.xml_object.bands.append(st_band)

    # Validate the XML
    espa_metadata.validate()

    # Write it to the XML file
    espa_metadata.write()


def write_output(xml_filename, espa_metadata, emis_b10_band_filename,
                 output_array, scale, offset, output_type, data_type):
    """Write the surface temperature band and its associated XML metadata entry.
       This can also write the ST Uncertainty band.       

    Args:
        xml_filename <str>: Filename for the ESPA Metadata XML
        espa_metadata <espa.Metadata>: XML metadata information
        emis_b10_band_filename <str>: template band filename
        output_array <numpy.2darray>: Array to write
        scale <float>: scale_factor for ST products
        offset <float>: add_offset for ST products
        output_type <str>: type of output, ST or UNC 
        data_type <str>: Data type field in metadata
    """

    # Set up output information.  Make it like the emissivity bands
    dataset = gdal.Open(emis_b10_band_filename)
    output_srs = osr.SpatialReference()
    output_srs.ImportFromWkt(dataset.GetProjection())
    output_transform = dataset.GetGeoTransform()
    samps = dataset.RasterXSize
    lines = dataset.RasterYSize
    del dataset
    product_id = espa_metadata.xml_object.global_metadata.product_id.text
    sensor_code = util.get_satellite_sensor_code(product_id)

    if output_type == "ST":
        # Write Surface Temperature band
        output_filename = ''.join([product_id, '_', sw_def.ST_BAND_NAME, '.img'])

        write_product(samps=samps,
                      lines=lines,
                      transform=output_transform,
                      wkt=output_srs.ExportToWkt(),
                      no_data_value=sw_def.L1_NO_DATA_VALUE,
                      data_type=data_type,
                      filename=output_filename,
                      image_data=output_array)

        # Write Surface Temperature band XML entry
        add_band_to_xml(espa_metadata=espa_metadata,
                        filename=output_filename,
                        product=sw_def.ST_PRODUCT,
                        name="surface_temperature",
                        category="image",
                        data_type="UINT16",
                        short_name='{0}ST'.format(sensor_code),
                        long_name="Surface Temperature",
                        data_units="temperature (kelvin)",
                        valid_range_min=str(sw_def.ST_RANGE_MIN),
                        valid_range_max=str(sw_def.ST_RANGE_MAX),
                        no_data_value=sw_def.L1_NO_DATA_VALUE,
                        scale=scale,
                        offset=offset)
    elif output_type == "UNC":
        output_filename = ''.join([product_id, '_', sw_def.STUNC_BAND_NAME, '.img'])

        write_product(samps=samps,
                      lines=lines,
                      transform=output_transform,
                      wkt=output_srs.ExportToWkt(),
                      no_data_value=sw_def.NO_DATA_VALUE,
                      data_type=data_type,
                      filename=output_filename,
                      image_data=output_array)

        # Write Surface Temperature uncertainty band XML entry
        add_band_to_xml(espa_metadata=espa_metadata,
                        filename=output_filename,
                        product=sw_def.STUNC_PRODUCT,
                        name=sw_def.STUNC_BAND_NAME,
                        category="qa",
                        data_type="INT16",
                        short_name='{0}STUNC'.format(sensor_code),
                        long_name="Surface temperature uncertainty band",
                        data_units="temperature (kelvin)",
                        valid_range_min=str(sw_def.STUNC_RANGE_MIN),
                        valid_range_max=str(sw_def.STUNC_RANGE_MAX),
                        no_data_value=sw_def.NO_DATA_VALUE,
                        scale=scale,
                        offset=offset)

    else: # Write TPW band

        output_filename = ''.join([product_id, output_type, '.img'])
        write_product(samps=samps,
                      lines=lines,
                      transform=output_transform,
                      wkt=output_srs.ExportToWkt(),
                      no_data_value=sw_def.NO_DATA_VALUE,
                      data_type=data_type,
                      filename=output_filename,
                      image_data=output_array)

        # Write TPW band XML entry.  This output is intended only for short term
        # internal use, so just borrow some ST_UNC definitions
        add_band_to_xml(espa_metadata=espa_metadata,
                        filename=output_filename,
                        product=sw_def.STUNC_PRODUCT,
                        name=output_type[1:], # Remove the leading "_"
                        category="qa",
                        data_type="FLOAT32",
                        short_name='{0}TPW'.format(sensor_code),
                        long_name="Total precipitable water band",
                        data_units="cm",
                        valid_range_min=str(0.0),
                        valid_range_max=str(sys.float_info.max),
                        no_data_value=sw_def.NO_DATA_VALUE,
                        scale=scale,
                        offset=offset)


def apply_sw_equation(bt_b10_array, bt_b11_array, emis_b10_array,
                      emis_b11_array, b):
    """Apply the core Split Window calculation to derive surface temperature.

    Args:
        bt_b10_array <numpy.2darray>: Brightness temperature for band 10
        bt_b11_array <numpy.2darray>: Brightness temperature for band 11
        emis_b10_array <numpy.2darray>: Emissivity band 10
        emis_b11_array <numpy.2darray>: Emissivity band 11
        b <[float]>: b-coefficients for Split Window equation
    """

    # Build kernel for convolution
    kernel = Box2DKernel(sw_def.KERNEL_SIZE)

    # Calculate the 5x5 averaged brightness temperature for Split Window
    # difference terms
    bt_b10_ave = ap_convolve(bt_b10_array, kernel, nan_treatment='interpolate',
                             normalize_kernel=True, boundary='extend')
    bt_b11_ave = ap_convolve(bt_b11_array, kernel, nan_treatment='interpolate',
                             normalize_kernel=True, boundary='extend')

    # Calculate terms for split window algorithm
    e_mean = (emis_b10_array + emis_b11_array) * 0.5
    e_diff = (1 - e_mean) / (e_mean)
    e_change = (emis_b10_array - emis_b11_array) / (e_mean**2)

    # Clean up memory
    del e_mean

    # Calculate terms for split window algorithm
    T_diff = (bt_b10_ave - bt_b11_ave) * 0.5
    T_plus = (bt_b10_array + bt_b11_array) * 0.5

    # Apply split window formula to derive surface temperature
    return b[0] + b[1] * T_plus + b[2] * T_plus * e_diff \
        + b[3] * T_plus * e_change + b[4] * T_diff \
        + b[5] * T_diff * e_diff + b[6] * T_diff * e_change + b[7] \
        * (bt_b10_ave - bt_b11_ave)**2


def apply_fill(surface_temperature, bt_b10_array, bt_b11_array, emis_b10_array,
               emis_b11_array, thermal_b10_array, thermal_b11_array):
    """Apply fill to ST in locations where other bands needed to compute ST are
       fill.

    Args:
        surface_temperature <numpy.2darray>: Surface Temperature band
        bt_b10_array <numpy.2darray>: Brightness temperature for band 10
        bt_b11_array <numpy.2darray>: Brightness temperature for band 11
        emis_b10_array <numpy.2darray>: Emissivity band 10
        emis_b11_array <numpy.2darray>: Emissivity band 11
        thermal_b10_array <str>: Thermal band 10
        thermal_b11_array <str>: Thermal band 11
    """

    # Find emissivity fill locations
    emis_fill_locations = np.where((emis_b10_array == sw_def.NO_DATA_VALUE) |
                                   (emis_b11_array == sw_def.NO_DATA_VALUE))

    # Find thermal fill locations
    thermal_fill_locations = \
        np.where((thermal_b10_array == sw_def.L1_NO_DATA_VALUE) |
                 (thermal_b11_array == sw_def.L1_NO_DATA_VALUE))

    # Set emissivity and thermal fill locations to the fill value in ST
    surface_temperature[emis_fill_locations] = sw_def.L1_NO_DATA_VALUE
    surface_temperature[thermal_fill_locations] = sw_def.L1_NO_DATA_VALUE

    return surface_temperature


def calculate_error(T10, T11, emis_band10_filename, emis_band11_filename,
                    emis_stdev_band10_filename, emis_stdev_band11_filename,
                    emis_source_filename, b, tpw, tpw_c1, tpw_c2, tpw_c3,
                    max_uncertainty, satellite, T10_error, T11_error,
                    ged_c_total_error_b10, ged_c_total_error_b11,
                    camel_c_total_error_b10, camel_c_total_error_b11,
                    ged_c1_val_b10, ged_c2_val_b10, ged_c1_val_b11,
                    ged_c2_val_b11, camel_c1_val_b10, camel_c2_val_b10,
                    camel_c1_val_b11, camel_c2_val_b11, ged_corr_emis,
                    camel_corr_emis_b10, camel_corr_emis_b11, corr_emis,
                    corr_appTemp, qa_pixel_band_filename, xml_filename,
                    espa_metadata):
    """ Calculate uncertainty metric by adding error in quadrature

    Args:
        bt_b10_array <numpy.2darray>: Brightness temperature for band 10
        bt_b11_array <numpy.2darray>: Brightness temperature for band 11
        emis_band10_filename <str>: Filename for emissivity band 10
        emis_band11_filename <str>: Filename for emissivity band 11
        emis_stdev_band10 filename <str>: Filename for band 10 emissivity stdev
        emis_stdev_band11_filename <str>: Filename for band 11 emissivity stdev
        emis_source_filename <str>: Filename for emissivity source
        b <[float]>: b-coefficients for Split Window equation
        tpw <numpy.2darray>: Total precipitable water band
        tpw_c1 <float>: Coefficient 1 for calculating TPW 
        tpw_c2 <float>: Coefficient 2 for calculating TPW
        tpw_c3 <float>: Coefficient 3 for calculating TPW 
        max_uncertainty <float>: Maximum allowed uncertainty value
        satellite <str>: LANDSAT_8 or LANDSAT_9
        T10_error <float>: Uncertainty in apparent temperature
        T11_error <float>: Uncertainty in apparent temperature
        ged_c_total_error_b10<float>: Regression error associated with
                                      conversion for ASTER GED
        ged_c_total_error_b11<float>: Regression error associated with
                                      conversion for ASTER GED
        camel_c_total_error_b10<float>: Regression error associated with
                                        conversion for CAMEL 
        camel_c_total_error_b11<float>: Regression error associated with
                                        conversion for CAMEL
        ged_c1_val_b10<float>: Emissivity conversion coefficient for ASTER GED
        ged_c2_val_b10<float>: Emissivity conversion coefficient for ASTER GED
        ged_c1_val_b11<float>: Emissivity conversion coefficient for ASTER GED
        ged_c2_val_b11<float>: Emissivity conversion coefficient for ASTER GED
        camel_c1_val_b10<float>: Emissivity conversion coefficient for CAMEL
        camel_c2_val_b10<float>: Emissivity conversion coefficient for CAMEL
        camel_c1_val_b11<float>: Emissivity conversion coefficient for CAMEL
        camel_c2_val_b11<float>: Emissivity conversion coefficient for CAMEL
        ged_corr_emis<float>: Cross-band correlation coefficient for emissivity
                              for ASTER GED
        camel_corr_emis_b10<float>: Cross-band correlation coefficient for
                                    emissivity for CAMEL b10
        camel_corr_emis_b11<float>: Cross-band correlation coefficient for 
                                    emissivity for CAMEL b11
        corr_emis<float>: Correlation coefficient for emissivity
        corr_appTemp<float>: Correlation coefficient for apparent temperature
        qa_pixel_band_filename <str>: Filename for pixel QA band
        xml_filename <str>: Filename for the ESPA Metadata XML
        espa_metadata <espa.Metadata>: XML metadata information
    """

    # Read emissivity standard deviation and source bands
    b10_std = Dataset.extract_raster_data(emis_stdev_band10_filename, 1)
    b11_std = Dataset.extract_raster_data(emis_stdev_band11_filename, 1)
    emis_source = Dataset.extract_raster_data(emis_source_filename, 1)

    # Split window coefficients
    b1 = b[1]
    b2 = b[2]
    b3 = b[3]
    b4 = b[4]
    b5 = b[5]
    b6 = b[6]
    b7 = b[7]

    # The updated approach uses MODTRAN profiles used for training the SW
    # and modelling the behaviour using a quadratic form.  In this approach
    # the b_total error is a function of TPW.  It uses MODTRAN to calculate
    # algorithmic error and embeds that in a b coefficient.
    b_total_error = tpw_c1 * (tpw**2) + tpw_c2 * tpw + tpw_c3

    # Clean up memory
    del tpw

    # Find locations derived from GED and CAMEL
    ged_locations = np.where(emis_source == cm_util.GED_EMIS_SOURCE)
    camel_locations = np.where(emis_source == cm_util.CAMEL_EMIS_SOURCE)

    # Calculate the uncertainty in the emissivity uncertainty calculation
    # with adding the covariance term (adding error in quadrature).  Water
    # and snow locations will remain at 0.
    e10_error_cov = np.zeros_like(emis_source, dtype=np.float32)
    e10_error_cov[ged_locations] = np.sqrt((ged_c_total_error_b10)**2 \
                  + (ged_c1_val_b10 * b10_std[ged_locations])**2 \
                  + (ged_c2_val_b10 * b11_std[ged_locations])**2 \
                  + 2 * ged_corr_emis * ged_c1_val_b10 * ged_c2_val_b10 \
                  * b10_std[ged_locations] * b11_std[ged_locations])
    e10_error_cov[camel_locations] = np.sqrt((camel_c_total_error_b10)**2 \
                  + (camel_c1_val_b10 * b10_std[camel_locations])**2 \
                  + (camel_c2_val_b10 * b11_std[camel_locations])**2 \
                  + 2 * camel_corr_emis_b10 * camel_c1_val_b10 \
                  * camel_c2_val_b10 * b10_std[camel_locations] \
                  * b11_std[camel_locations])

    # Calculate the uncertainty in the emissivity uncertainty calculation
    # with adding the covariance term (adding error in quadrature).  Water
    # and snow locations will remain at 0.
    e11_error_cov = np.zeros_like(emis_source, dtype=np.float32)

    # Clean up memory
    del emis_source

    e11_error_cov[ged_locations] = np.sqrt((ged_c_total_error_b11)**2 \
                  + (ged_c1_val_b11 * b10_std[ged_locations])**2 \
                  + (ged_c2_val_b11 * b11_std[ged_locations])**2 \
                  + 2 * ged_corr_emis * ged_c1_val_b11 * ged_c2_val_b11 \
                  * b10_std[ged_locations] * b11_std[ged_locations])
    e11_error_cov[camel_locations] = np.sqrt((camel_c_total_error_b11)**2 \
                  + (camel_c1_val_b11 * b10_std[camel_locations])**2 \
                  + (camel_c2_val_b11 * b11_std[camel_locations])**2 \
                  + 2 * camel_corr_emis_b11 * camel_c1_val_b11 \
                  * camel_c2_val_b11 * b10_std[camel_locations] \
                  * b11_std[camel_locations])

    # Clean up memory
    del b10_std
    del b11_std

    # Read emissivity bands
    e10 = Dataset.extract_raster_data(emis_band10_filename, 1)
    e11 = Dataset.extract_raster_data(emis_band11_filename, 1)

    # Partial derivatives of SW algorithm (see prototype for how this was
    # calculated)
    T10diff = b1/2 + b2 * (-e10/2 - e11/2 + 1)/(2 * (e10/2 + e11/2)) \
            + b3 * (e10 - e11)/(2 * (e10/2 + e11/2)**2) \
            + b4/2 + b5 * (-e10/2 - e11/2 + 1)/(2 * (e10/2 + e11/2)) \
            + b6 * (e10 - e11)/(2 * (e10/2 + e11/2)**2) \
            + b7 * (2 * T10 - 2 * T11)
    T11diff = b1/2 + b2 * (-e10/2 - e11/2 + 1)/(2 * (e10/2 + e11/2)) \
            + b3 * (e10 - e11)/(2 * (e10/2 + e11/2)**2) \
            - b4/2 - b5 * (-e10/2 - e11/2 + 1)/(2 * (e10/2 + e11/2)) \
            - b6 * (e10 - e11)/(2 * (e10/2 + e11/2)**2) \
            + b7 * (-2 * T10 + 2 * T11)
    E10diff = -b2 * (T10 + T11)/(4 * (e10/2 + e11/2)) \
            - b2 * (T10 + T11) * (-e10/2 - e11/2 + 1)/(4 * (e10/2 + e11/2)**2) \
            + b3 * (T10 + T11)/(2 * (e10/2 + e11/2)**2) \
            - b3 * (T10 + T11) * (e10 - e11)/(2 * (e10/2 + e11/2)**3) \
            - b5 * (T10 - T11)/(4 * (e10/2 + e11/2)) \
            - b5 * (T10 - T11) * (-e10/2 - e11/2 + 1)/(4 * (e10/2 + e11/2)**2) \
            + b6 * (T10 - T11)/(2 * (e10/2 + e11/2)**2) \
            - b6 * (T10 - T11) * (e10 - e11)/(2 * (e10/2 + e11/2)**3)
    E11diff = -b2 * (T10 + T11)/(4 * (e10/2 + e11/2)) \
            - b2 * (T10 + T11) * (-e10/2 - e11/2 + 1)/(4 * (e10/2 + e11/2)**2) \
            - b3 * (T10 + T11)/(2 * (e10/2 + e11/2)**2) \
            - b3 * (T10 + T11) * (e10 - e11)/(2 * (e10/2 + e11/2)**3) \
            - b5 * (T10 - T11)/(4 * (e10/2 + e11/2)) \
            - b5 * (T10 - T11) * (-e10/2 - e11/2 + 1)/(4 * (e10/2 + e11/2)**2) \
            - b6 * (T10 - T11)/(2 * (e10/2 + e11/2)**2) \
            - b6 * (T10 - T11) * (e10 - e11)/(2 * (e10/2 + e11/2)**3)

    # Clean up memory
    del T10
    del T11
    del e10
    del e11

    # Split Window uncertainty in quadrature with and without correlation
    error_cov = np.sqrt(b_total_error**2 + (T10diff * T10_error)**2 \
                + (T11diff * T11_error)**2 + (E10diff * e10_error_cov)**2 \
                + (E11diff * e11_error_cov)**2 \
                + 2 * corr_appTemp * T10diff * T11diff * T10_error * T11_error \
                + 2 * corr_emis * E10diff * E11diff * e10_error_cov \
                * e11_error_cov)

    # Get QA pixel mask for clouds
    cloud_mask = get_cloud_mask(qa_pixel_band_filename)

    # Apply cloud mask
    error_cov[cloud_mask == 0] = sw_def.NO_DATA_VALUE

    # Correct for artifacts caused by the ASTER data
    error_cov[error_cov > max_uncertainty] = max_uncertainty

    return error_cov


def get_radiance(thermal_b10_array, thermal_b11_array, thermal_b10_scale,
                 thermal_b11_scale, thermal_b10_offset, thermal_b11_offset):
    """ Build radiance arrays suitable as input to the total precipitable water
        procedure

    Args:
        thermal_b10_array <numpy.2darray>: Thermal band 10 data
        thermal_b11_array <numpy.2darray>: Thermal band 11 data
        thermal_b10_scale <float>: scale_factor for thermal band 10
        thermal_b11_scale <float>: scale_factor for thermal band 11
        thermal_b10_offset <float>: add_offset for thermal band 10
        thermal_b11_offset <float>: add_offset for thermal band 11

    Returns:
        b10_radiance_array <numpy.2darray>: Band 10 radiance 
        b11_radiance_array <numpy.2darray>: Band 11 radiance 
    """

    # Convert thermal bands DN to radiance
    b10_radiance_array = thermal_b10_scale * thermal_b10_array \
                                           + thermal_b10_offset
    b11_radiance_array = thermal_b11_scale * thermal_b11_array \
                                           + thermal_b11_offset
    b10_radiance_array[b10_radiance_array == thermal_b10_offset] \
                       = sw_def.TPW_NO_DATA_VALUE 
    b11_radiance_array[b11_radiance_array == thermal_b11_offset] \
                       = sw_def.TPW_NO_DATA_VALUE 

    # Set radiance bands to TPW fill value where either thermal band was fill

    # Get mask locations from the thermal bands
    b10_mask = (thermal_b10_array == sw_def.L1_NO_DATA_VALUE)
    b11_mask = (thermal_b11_array == sw_def.L1_NO_DATA_VALUE)

    # Combine the mask locations from the thermal bands into a single mask
    combined_mask = np.logical_or(b11_mask, b10_mask)

    # Set the mask locations to the TPW fill value
    b10_radiance_array[combined_mask] = sw_def.TPW_NO_DATA_VALUE
    b11_radiance_array[combined_mask] = sw_def.TPW_NO_DATA_VALUE

    # Clean up memory
    del b10_mask
    del b11_mask
    del combined_mask
    del thermal_b10_array
    del thermal_b11_array

    b10_radiance_array = b10_radiance_array.astype('float32')
    b11_radiance_array = b11_radiance_array.astype('float32')

    return (b10_radiance_array, b11_radiance_array)


def get_cloud_mask(qa_pixel_band_filename):
    """ Get a mask made by combining several Pixel QA cloud bits.

    Args:
        qa_pixel_band_filename <str>: Filename for the QA Pixel band

    Returns:
        cloud_mask <numpy.2darray>: Cloud mask band
    """

    qa_pixel_array = Dataset.extract_raster_data(qa_pixel_band_filename, 1)
    return ((qa_pixel_array & (sw_def.PQA_DILATED_CLOUD | \
             sw_def.PQA_CIRRUS | sw_def.PQA_CLOUD)) == 0)


def split_window(xml_filename, data_path, scale, offset):
    """Provides the main processing algorithm for generating the ST product
       using the Split Window algorithm.

    Args:
        xml_filename <str>: Filename for the ESPA Metadata XML
        data_path <str>: Directory for Surface Temperature data files
        scale <float>: scale_factor for ST products
        offset <float>: add_offset for ST products
    """

    # XML Metadata
    espa_metadata = Metadata(xml_filename)
    espa_metadata.parse()

    # Determine if it's a TIRS-only scenario
    product_id = espa_metadata.xml_object.global_metadata.product_id.text
    sensor_code = util.get_satellite_sensor_code(product_id)
    tirs_only = False
    if sensor_code in ['LT08', 'LT09']:
        tirs_only = True

    # Read metadata
    (emis_b10_band_filename, emis_b11_band_filename,
     emis_stdev_b10_band_filename, emis_stdev_b11_band_filename,
     emis_source_band_filename, qa_pixel_band_filename,
     thermal_b10_band_filename, thermal_b11_band_filename, thermal_b10_scale,
     thermal_b10_offset, thermal_b11_scale, thermal_b11_offset) \
         = (read_metadata(espa_metadata))

    # Read thermal bands
    thermal_b10_array = Dataset.extract_raster_data(thermal_b10_band_filename, 1)
    thermal_b11_array = Dataset.extract_raster_data(thermal_b11_band_filename, 1)

    # Build radiance arrays used as input to TPW process
    (b10_radiance_array, b11_radiance_array) = \
            get_radiance(thermal_b10_array, thermal_b11_array,
                         thermal_b10_scale, thermal_b11_scale,
                         thermal_b10_offset, thermal_b11_offset)

    # Clean memory
    del thermal_b11_array

    # Get total precipitable water
    model_dir = os.path.join(data_path, sw_def.XGBOOST_DIR)
    tpw = pred_tpw(model_dir, b10_radiance_array, b11_radiance_array)

    # Use the masked B10 radiance fill locations to set TPW fill
    tpw[thermal_b10_array == sw_def.L1_NO_DATA_VALUE] \
        = sw_def.L1_NO_DATA_VALUE

    # Clean memory
    del thermal_b10_array

    # Read emissivity bands
    emis_b10_array = Dataset.extract_raster_data(emis_b10_band_filename, 1)
    emis_b11_array = Dataset.extract_raster_data(emis_b11_band_filename, 1)

    # Read b-coefficients.  These were derived  as per journal article: Towards
    # an Operational, Split Window-Derived Surface Temperature Product for
    # the Thermal Infrared Sensors Onboard Landsat 8 and 9, A.Gerace,
    # T.Kleynhans, R.Eon, M. Montanaro
    satellite = espa_metadata.xml_object.global_metadata.satellite
    coefficients_filename = sw_def.COEFFICIENTS_FILENAME

    coefficient_file = os.path.join(data_path, coefficients_filename)
    with open(coefficient_file) as coefficient_data:
        coeff = json.load(coefficient_data)
        coefficient_data.close()
    sw_coeff = coeff['Split_Window']

    # T10/T11_error: Uncertainty in apparent temperature - hardcoded based on 
    # Montanaro et al.: Derivation and validation of the stray light correction 
    # algorithm for the thermal infrared sensor onboard Landsat 8 and as
    # discussed in virtual CalVal meeting 2020 March
    if satellite == 'LANDSAT_8':
        b=sw_coeff['L8']['SW_coeff']
        T10_error=sw_coeff['L8_B10']['landsat_uncertainty']
        T11_error=sw_coeff['L8_B11']['landsat_uncertainty']
        ged_c_total_error_b10=sw_coeff['L8_B10']['emis_regfit']
        ged_c_total_error_b11=sw_coeff['L8_B11']['emis_regfit']
        camel_c_total_error_b10=sw_coeff['L8_B10']['emis_regfit_camel']
        camel_c_total_error_b11=sw_coeff['L8_B11']['emis_regfit_camel']
        ged_c1_val_b10=sw_coeff['L8_B10']['estimated_1']
        ged_c2_val_b10=sw_coeff['L8_B10']['estimated_2']
        ged_c1_val_b11=sw_coeff['L8_B11']['estimated_1']
        ged_c2_val_b11=sw_coeff['L8_B11']['estimated_2']
        camel_c1_val_b10=sw_coeff['L8_B10']['camel_coeff_1']
        camel_c2_val_b10=sw_coeff['L8_B10']['camel_coeff_2']
        camel_c1_val_b11=sw_coeff['L8_B11']['camel_coeff_1']
        camel_c2_val_b11=sw_coeff['L8_B11']['camel_coeff_2']
    else:
        b=sw_coeff['L9']['SW_coeff']
        T10_error=sw_coeff['L9_B10']['landsat_uncertainty']
        T11_error=sw_coeff['L9_B11']['landsat_uncertainty']
        ged_c_total_error_b10=sw_coeff['L9_B10']['emis_regfit']
        ged_c_total_error_b11=sw_coeff['L9_B11']['emis_regfit']
        camel_c_total_error_b10=sw_coeff['L9_B10']['emis_regfit_camel']
        camel_c_total_error_b11=sw_coeff['L9_B11']['emis_regfit_camel']
        ged_c1_val_b10=sw_coeff['L9_B10']['estimated_1']
        ged_c2_val_b10=sw_coeff['L9_B10']['estimated_2']
        ged_c1_val_b11=sw_coeff['L9_B11']['estimated_1']
        ged_c2_val_b11=sw_coeff['L9_B11']['estimated_2']
        camel_c1_val_b10=sw_coeff['L9_B10']['camel_coeff_1']
        camel_c2_val_b10=sw_coeff['L9_B10']['camel_coeff_2']
        camel_c1_val_b11=sw_coeff['L9_B11']['camel_coeff_1']
        camel_c2_val_b11=sw_coeff['L9_B11']['camel_coeff_2']
    tpw_c1=sw_coeff['tpw']['c1']
    tpw_c2=sw_coeff['tpw']['c2']
    tpw_c3=sw_coeff['tpw']['c3']
    # Correlation coefficient for apparent temperature and emissivity
    # This was calculated using the 113 MODIS emissivities (spectrally
    # sampled)
    # The appTemp (apparent temperature between B10 and B11) correlation
    # was calculated using the TIGR simulation data
    corr_emis=sw_coeff['corr_emis']
    ged_corr_emis=sw_coeff['corr_emis_aster']
    camel_corr_emis_b10=sw_coeff['corr_camel_9_11']
    camel_corr_emis_b11=sw_coeff['corr_camel_11_12']
    corr_app_temp=sw_coeff['corr_app_temp']
    max_uncertainty=sw_coeff['max_uncertainty']

    if len(b) != 8:
        raise InvalidCoefficientsError('Incorrect number of coefficients' +
                                       ' read from:' + coefficient_file)

    # Compute BT from LUT.  This is more accurate than using K coefficient
    # formula.
    conversion_lut_name = util.get_radiometric_conversion_filename(satellite)

    # For L8-9, the columns are:
    #
    # 1. Temperature
    # 2. Radiance in RTTOV units (wavenumber space) for band 10
    # 3. Radiance in RTTOV units (wavenumber space) for band 11
    # 4. Radiance in microns for band 10
    # 5. Radiance in microns for band 11
    #
    # In this phase we only use the temperature and radiance in microns
    temperature_col = 0
    b10_radiance_col = 3
    b11_radiance_col = 4
    conversion_data = np.loadtxt(os.path.join(data_path, conversion_lut_name),
                     dtype=float, delimiter=' ', usecols=(temperature_col,
                     b10_radiance_col, b11_radiance_col))
    temp_lut = conversion_data[:, 0]         # Temperature
    b10_radiance_lut = conversion_data[:, 1] # B10 radiance in microns
    b11_radiance_lut = conversion_data[:, 2] # B11 radiance in microns

    # Compute brightness temperature, applying the LUT and interpolation
    # to the radiance bands
    bt_b10_array = np.interp(b10_radiance_array, b10_radiance_lut, temp_lut)
    bt_b11_array = np.interp(b11_radiance_array, b11_radiance_lut, temp_lut)
    bt_b10_array[tpw == sw_def.L1_NO_DATA_VALUE] = sw_def.L1_NO_DATA_VALUE
    bt_b11_array[tpw == sw_def.L1_NO_DATA_VALUE] = sw_def.L1_NO_DATA_VALUE

    # Clean memory
    del b10_radiance_array
    del b11_radiance_array

    # Brightness temperature arrays must use float for np.nan
    bt_b10_array = bt_b10_array.astype(float)
    bt_b11_array = bt_b11_array.astype(float)

    # Brightness temperature array fill must be nan for astropy function to work
    bt_b10_array[bt_b10_array == sw_def.L1_NO_DATA_VALUE] = np.nan
    bt_b11_array[bt_b11_array == sw_def.L1_NO_DATA_VALUE] = np.nan

    # Apply the Split Window equation to derive Surface Temperature
    surface_temperature = apply_sw_equation(bt_b10_array, bt_b11_array,
                                            emis_b10_array, emis_b11_array, b)

    # Read thermal bands
    thermal_b10_array = Dataset.extract_raster_data(thermal_b10_band_filename, 1)
    thermal_b11_array = Dataset.extract_raster_data(thermal_b11_band_filename, 1)

    # Apply fill to the Surface Temperature band
    surface_temperature = apply_fill(surface_temperature, bt_b10_array,
                                     bt_b11_array, emis_b10_array,
                                     emis_b11_array, thermal_b10_array,
                                     thermal_b11_array)

    # Find thermal high saturation locations
    thermal_saturation_locations = \
        np.where((thermal_b10_array == util.OLI_HIGH_SATURATION) |
                 (thermal_b11_array == util.OLI_HIGH_SATURATION))

    # Clean memory
    del emis_b10_array
    del emis_b11_array
    del thermal_b10_array
    del thermal_b11_array

    # Scale the data
    surface_temperature[surface_temperature != sw_def.L1_NO_DATA_VALUE] -= \
        offset
    surface_temperature[surface_temperature != sw_def.L1_NO_DATA_VALUE] /= \
        scale

    # If the result is outside the new range, put back in the range
    surface_temperature[surface_temperature > sw_def.MAX_UINT16] = \
        sw_def.MAX_UINT16
    surface_temperature[surface_temperature < sw_def.MIN_UINT16] = \
        sw_def.MIN_UINT16

    # Remove input saturation locations in the Surface Temperature band
    surface_temperature[thermal_saturation_locations] = sw_def.L1_NO_DATA_VALUE

    # Write surface temperature output band and associated metadata
    write_output(xml_filename, espa_metadata, emis_b10_band_filename,
                 surface_temperature, scale, offset, "ST", gdal.GDT_UInt16)

    # Write TPW output band and associated metadata
    write_output(xml_filename, espa_metadata, emis_b10_band_filename,
                 tpw, 1, 0, "_tpw", gdal.GDT_Float32)

    # Reduce band size from float64 to float32 to save memory
    bt_b10_array = bt_b10_array.astype(np.float32)
    bt_b11_array = bt_b11_array.astype(np.float32)

    split_window_UNC = calculate_error(bt_b10_array, bt_b11_array,
                                       emis_b10_band_filename,
                                       emis_b11_band_filename,
                                       emis_stdev_b10_band_filename,
                                       emis_stdev_b11_band_filename,
                                       emis_source_band_filename, b, tpw,
                                       tpw_c1, tpw_c2, tpw_c3, max_uncertainty,
                                       satellite, T10_error, T11_error,
                                       ged_c_total_error_b10,
                                       ged_c_total_error_b11,
                                       camel_c_total_error_b10,
                                       camel_c_total_error_b11,
                                       ged_c1_val_b10, ged_c2_val_b10,
                                       ged_c1_val_b11, ged_c2_val_b11,
                                       camel_c1_val_b10, camel_c2_val_b10,
                                       camel_c1_val_b11, camel_c2_val_b11,
                                       ged_corr_emis, camel_corr_emis_b10,
                                       camel_corr_emis_b11, corr_emis,
                                       corr_app_temp, qa_pixel_band_filename,
                                       xml_filename, espa_metadata)

    # Apply fill to the Surface Temperature UNC band
    split_window_UNC[surface_temperature == sw_def.L1_NO_DATA_VALUE] = \
        sw_def.NO_DATA_VALUE

    # Clean up memory
    del surface_temperature

    # Scale the UNC band data
    split_window_UNC[split_window_UNC != sw_def.NO_DATA_VALUE] -= \
        sw_def.STUNC_OFFSET
    split_window_UNC[split_window_UNC != sw_def.NO_DATA_VALUE] /= \
        sw_def.STUNC_SCALE

    # If the result is outside the new range, put back in the range.
    # We wouldn't expect negative values except nodata
    split_window_UNC[split_window_UNC > sw_def.MAX_INT16] = \
        sw_def.MAX_INT16
    split_window_UNC[split_window_UNC < sw_def.MIN_INT16] = \
        sw_def.MIN_INT16

    # Remove input saturation locations in the Surface Temperature band
    split_window_UNC[thermal_saturation_locations] = sw_def.NO_DATA_VALUE

    # Clean memory
    del thermal_saturation_locations

    # Write surface temperature UNC output band and associated metadata
    write_output(xml_filename, espa_metadata, emis_b10_band_filename,
                 split_window_UNC, sw_def.STUNC_SCALE, sw_def.STUNC_OFFSET, "UNC",
                 gdal.GDT_Int16)


def main():
    """Main processing for running the Split Window ST algorithm
    """

    # Command Line Arguments
    args = retrieve_command_line_arguments()

    # Configure logging
    Sys.setup_logging()

    # Set GDAL exceptions
    gdal.UseExceptions()

    logger.info('*** Begin ST Split Window algorithm ***')

    try:
        split_window(xml_filename=args.xml_filename,
                     data_path=args.data_path,
                     scale=args.scale,
                     offset=args.offset)

    except Exception:
        logger.exception('Processing failed')
        sys.exit(1)  # EXIT FAILURE

    logger.info('*** Run ST Split Window algorithm - Complete ***')


if __name__ == '__main__':
    main()
