#! /usr/bin/env python

'''
    File: st_generate_qa.py

    Purpose: Builds an uncertainty band for the surface temperature product.

    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
import re
from argparse import ArgumentParser
from collections import namedtuple

import numpy as np
from lxml import objectify as objectify
from osgeo import gdal, osr

import st_utilities as util
from st_exceptions import MissingBandError
from espa import Metadata, Sys, Geo

logger = logging.getLogger(__name__)
SourceInfo = namedtuple('SourceInfo', ('filename'))
ThermalConstantInfo = namedtuple('ThermalInfo', ('k1', 'k2'))
CoefficientInfo = namedtuple('CoefficientInfo',
                             ('transmittance_coeff_1',
                              'transmittance_coeff_2',
                              'transmittance_coeff_3',
                              'transmittance_lower_bound',
                              'upwelled_radiance_coeff_1',
                              'upwelled_radiance_coeff_2',
                              'upwelled_radiance_coeff_3',
                              'upwelled_radiance_upper_bound',
                              'downwelled_radiance_coeff_1',
                              'downwelled_radiance_coeff_2',
                              'downwelled_radiance_coeff_3',
                              'downwelled_radiance_upper_bound'))


# send floating-point error messages to stdout for capture in log
np.seterr(all='print')


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

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

    parser = ArgumentParser(description='Builds surface temperature UNC band')

    parser.add_argument('--version',
                        action='version',
                        version=util.Version.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('--num_threads',
                        action='store', dest='num_threads',
                        required=False, default=1,
                        help='Number of threads to use if OpenMP enabled')

    args = parser.parse_args()

    # Verify that the --xml parameter was specified
    if args.xml_filename is None:
        raise Exception('--xml must be specified on the command line')

    return args


def retrieve_metadata_information(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:
        <SourceInfo>: Populated with source information
    """

    intermediate_filename = None

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

    # Error if we didn't find the required intermediate band in the data
    if intermediate_filename is None:
        raise MissingBandError('Failed to find the intermediate band {0}'
                               ' in the input data'.format(band_name))

    return SourceInfo(filename=intermediate_filename)


def retrieve_thermal_constants(espa_metadata, satellite):
    """Reads thermal constants from the metadata XML file

    Args:
        espa_metadata <espa.Metadata>: XML metadata
        satellite <str>: Name of satellite

    Returns:
        <ThermalConstantInfo>: Populated with thermal constants
    """

    thermal_constants = None
    band_name = None

    # Determine the band to retrieve constants from using the satellite
    if satellite == 'LANDSAT_4' or satellite == 'LANDSAT_5':
        band_name = "b6"
    elif satellite == 'LANDSAT_7':
        band_name = "b61"
    elif satellite == 'LANDSAT_8' or satellite == 'LANDSAT_9':
        band_name = "b10"
    else:
        raise Exception('Unsupported satellite')

    # Find the band to extract information from
    for band in espa_metadata.xml_object.bands.band:
        if (band.get('name') == band_name):
            thermal_constants = str(band.thermal_const)
            k1 = band.thermal_const.get('k1')
            k2 = band.thermal_const.get('k2')

    # Error if we didn't find the required intermediate band in the data
    if thermal_constants is None:
        raise MissingBandError('Failed to find the thermal band in the '
                               ' input data')

    return ThermalConstantInfo(k1=k1, k2=k2)


def retrieve_coefficients(satellite, st_data_dir):
    """Reads coefficients from JSON coefficient file

    These updated values were provided by Rehman Eon from RIT via email
    correspondence related to a Cal/Val TIM on 6/12/2025.

    Args:
        satellite <str>: Name of satellite
        st_data_dir <str>: Location of the ST data files

    Returns:
        <CoefficientInfo>: Populated coefficients
    """

    # Read coefficient file.
    coefficient_file = os.path.join(st_data_dir, util.COEFFICIENT_FILE)
    with open(coefficient_file) as coefficient_data:
        all_coeff = json.load(coefficient_data)
        coefficient_data.close()

    if satellite == 'LANDSAT_4':
        coeff = all_coeff['Single_Channel']['L4']
    elif satellite == 'LANDSAT_5':
        coeff = all_coeff['Single_Channel']['L5']
    elif satellite == 'LANDSAT_7':
        coeff = all_coeff['Single_Channel']['L7']
    elif satellite == 'LANDSAT_8': # We only process band10
        coeff = all_coeff['Single_Channel']['L8_B10']
    elif satellite == 'LANDSAT_9': # We only process band10
        coeff = all_coeff['Single_Channel']['L9_B10']
    else:
        raise Exception('Unsupported satellite')

    return CoefficientInfo(
        transmittance_coeff_1=coeff['transmittance_coeff_1'],
        transmittance_coeff_2=coeff['transmittance_coeff_2'],
        transmittance_coeff_3=coeff['transmittance_coeff_3'],
        transmittance_lower_bound=coeff['transmittance_lower_bound'],
        upwelled_radiance_coeff_1=coeff['upwelled_radiance_coeff_1'],
        upwelled_radiance_coeff_2=coeff['upwelled_radiance_coeff_2'],
        upwelled_radiance_coeff_3=coeff['upwelled_radiance_coeff_3'],
        upwelled_radiance_upper_bound=coeff['upwelled_radiance_upper_bound'],
        downwelled_radiance_coeff_1=coeff['downwelled_radiance_coeff_1'],
        downwelled_radiance_coeff_2=coeff['downwelled_radiance_coeff_2'],
        downwelled_radiance_coeff_3=coeff['downwelled_radiance_coeff_3'],
        downwelled_radiance_upper_bound=coeff['downwelled_radiance_upper_bound'])


def calculate_unc(unc_img_filename, radiance_filename, transmission_filename,
                  upwelled_filename, downwelled_filename, emis_filename,
                  emis_stdev_filename, thermal_1_filename, thermal_2_filename,
                  coefficients, satellite, k1, k2, fill_value, lines, samps,
                  num_threads):
    """Calculate UNC

    Args:
        unc_img_filename <str>: Name of the output uncertainty file
        radiance_filename <str>: Name of radiance file
        transmission_filename <str>: Name of atmospheric transmission file
        upwelled_filename <str>: Name of upwelled radiance file
        downwelled_filename <str>: Name of downwelled radiance file
        emis_filename <str>: Name of emissivity file
        emis_stdev_filename <str>: Name of emissivity standard deviation file
        thermal_1_filename <str>: Name of level 1 thermal band 1
        thermal_2_filename <str>: Name of level 1 thermal band 2 (ETM)
        coefficients <CoefficientInfo>: satellite-specific coefficients values
        satellite <str>: Name of satellite (e.g.: "LANDSAT_8")
        k1 <float>: K1 thermal conversion constant for the satellite
        k2 <float>: K2 thermal conversion constant for the satellite
        fill_value <float>: No data (fill) value to use
        lines <int>: Number of lines in band
        samps <int>: Number of samples in band
        num_threads <int>: Number of processing threads to use

    Returns:
        <numpy.2darray>: Generated surface temperature UNC band data
    """

    logger.info('Building UNC band')

    cmd = ['calculate_surface_temp_unc', str(satellite), str(lines), str(samps),
           radiance_filename, transmission_filename, upwelled_filename,
           downwelled_filename, emis_filename, emis_stdev_filename,
           thermal_1_filename, thermal_2_filename, unc_img_filename, str(k1),
           str(k2), str(coefficients.transmittance_coeff_1),
           str(coefficients.transmittance_coeff_2),
           str(coefficients.transmittance_coeff_3),
           str(coefficients.transmittance_lower_bound),
           str(coefficients.upwelled_radiance_coeff_1),
           str(coefficients.upwelled_radiance_coeff_2),
           str(coefficients.upwelled_radiance_coeff_3),
           str(coefficients.upwelled_radiance_upper_bound),
           str(coefficients.downwelled_radiance_coeff_1),
           str(coefficients.downwelled_radiance_coeff_2),
           str(coefficients.downwelled_radiance_coeff_3),
           str(coefficients.downwelled_radiance_upper_bound),
           str(fill_value),str(MULT_FACTOR), str(num_threads)]
           
    cmd = ' '.join(cmd)
    output = ''
    try:
        logger.info('Calling [%s]', cmd)
        output = Sys.execute_cmd(cmd)
    except Exception:
        logger.error('Failed creating uncertainty data')
        raise
    finally:
        if output: # Check if output is empty
            logger.info(output)


def add_unc_band_to_xml(espa_metadata, filename, sensor_code, no_data_value):
    """Adds the UNC band to the Metadata XML file

    Args:
        espa_metadata <espa.Metadata>: XML metadata information
        filename <str>: Full path for the output file to create
        sensor_code <str>: Name prefix for the sensor
        no_data_value <float>: Value to use for fill
    """

    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)

    unc_band = maker.band()
    unc_band.set('product', 'st_unc')
    unc_band.set('source', source_product)
    unc_band.set('name', 'st_unc')
    unc_band.set('category', 'qa')
    unc_band.set('data_type', 'INT16')
    unc_band.set('scale_factor', str(SCALE_FACTOR))
    unc_band.set('nlines', base_band.attrib['nlines'])
    unc_band.set('nsamps', base_band.attrib['nsamps'])
    unc_band.set('fill_value', str(no_data_value))

    unc_band.short_name = maker.element('{0}STUNC'.format(sensor_code))

    unc_band.long_name = maker.element('Surface temperature uncertainty band')
    unc_band.file_name = maker.element(filename)

    unc_band.pixel_size = base_band.pixel_size

    unc_band.resample_method = maker.element('none')
    unc_band.data_units = maker.element('temperature (kelvin)')

    unc_band.valid_range = maker.element()
    unc_band.valid_range.set('min', '0')
    unc_band.valid_range.set('max', '32767')

    unc_band.auxiliary_source = maker.element()
    unc_band.auxiliary_source.set('reanalysis', util.REANALYSIS.get_reanalysis())

    unc_band.app_version = maker.element(util.Version.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')))
    unc_band.production_date = maker.element(date_now)

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

    # Validate the XML
    espa_metadata.validate()

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


def write_unc_header(base_hdr, new_hdr):
    """Write an ENVI header based on an existing ENVI header.

    Args:
        base_hdr <str>: Base header file
        new_hdr <str>: New header file
    """

    # Read each line of the base header and write out to the new
    # header, updating information along the way.
    with open(base_hdr) as base_fp, open(new_hdr, 'w') as new_fp:
        for line in base_fp.readlines():
            if re.search("data type =", line):
                line = "data type = 2\n"
            # The band names value may be split across multiple lines.
            # Only update the name if it's on a single line.
            # (It's not a critical parameter, so not a big deal.)
            elif re.search("band names = {.+}", line):
                line = "band names = {st_unc}\n"
            new_fp.write(line)


def generate_unc(xml_filename, st_data_dir, no_data_value, num_threads):
    """Provides the main processing algorithm for generating the UNC product.

    Args:
        xml_filename <str>: Filename for the ESPA Metadata XML
        st_data_dir <str>: Location of the ST data files
        no_data_value <float>: No data (fill) value to use
        num_threads <int>: Number of processing threads to use
    """

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

    radiance_src_info \
        = retrieve_metadata_information(espa_metadata,
                                        'st_thermal_radiance')
    transmission_src_info \
        = retrieve_metadata_information(espa_metadata,
                                        'st_atmospheric_transmittance')
    upwelled_src_info \
        = retrieve_metadata_information(espa_metadata,
                                        'st_upwelled_radiance')
    downwelled_src_info \
        = retrieve_metadata_information(espa_metadata,
                                        'st_downwelled_radiance')
    emis_src_info = retrieve_metadata_information(espa_metadata, 'emis')
    emis_stdev_src_info \
        = retrieve_metadata_information(espa_metadata, 'emis_stdev')
    satellite = espa_metadata.xml_object.global_metadata.satellite

    # Determine the band to retrieve constants from using the satellite
    if satellite == 'LANDSAT_4' or satellite == 'LANDSAT_5':
        thermal_1_src_info = retrieve_metadata_information(espa_metadata, 'b6')
        thermal_2_src_info = SourceInfo('N/A') 
    elif satellite == 'LANDSAT_7':
        thermal_1_src_info = retrieve_metadata_information(espa_metadata, 'b61')
        thermal_2_src_info = retrieve_metadata_information(espa_metadata, 'b62')
    elif satellite == 'LANDSAT_8' or satellite == 'LANDSAT_9':
        thermal_1_src_info = retrieve_metadata_information(espa_metadata, 'b10')
        thermal_2_src_info = SourceInfo('N/A') 
    else:
        raise Exception('Unsupported satellite')
    thermal_info = retrieve_thermal_constants(espa_metadata, satellite)

    # Retrieve coefficients
    coefficients = retrieve_coefficients(satellite, st_data_dir)

    # Determine output information.  Make it like the emissivity band
    dataset = gdal.Open(emis_src_info.filename)
    output_srs = osr.SpatialReference()
    output_srs.ImportFromWkt(dataset.GetProjection())
    output_transform = dataset.GetGeoTransform()
    samps = dataset.RasterXSize
    lines = dataset.RasterYSize
    del dataset

    # Retrieve and validate sensor code
    product_id = espa_metadata.xml_object.global_metadata.product_id.text
    sensor_code = util.get_satellite_sensor_code(product_id)

    # Build UNC filename
    unc_img_filename = ''.join([product_id, '_st_unc', '.img'])

    # Build UNC information in memory
    calculate_unc(unc_img_filename,
                  radiance_src_info.filename,
                  transmission_src_info.filename,
                  upwelled_src_info.filename,
                  downwelled_src_info.filename,
                  emis_src_info.filename,
                  emis_stdev_src_info.filename,
                  thermal_1_src_info.filename,
                  thermal_2_src_info.filename,
                  coefficients,
                  satellite,
                  float(thermal_info.k1),
                  float(thermal_info.k2),
                  no_data_value, lines, samps, num_threads)

    # Write the UNC header file, based on the thermal 1 file.
    write_unc_header(thermal_1_src_info.filename.replace('.img', '.hdr'),
                     unc_img_filename.replace('.img', '.hdr'))

    add_unc_band_to_xml(espa_metadata=espa_metadata,
                        filename=unc_img_filename,
                        sensor_code=sensor_code,
                        no_data_value=no_data_value)


# Specify the no data, scale factor, and multiplication factor.
SCALE_FACTOR = 0.01
MULT_FACTOR = 100.0


def main():
    """Main processing for building the surface temperature UNC band
    """

    # Command Line Arguments
    args = retrieve_command_line_arguments()

    # Configure logging
    Sys.setup_logging()

    logger.info('*** Begin ST Generate UNC ***')

    try:
        # Set GDAL exceptions
        gdal.UseExceptions()

        # Register all the gdal drivers
        gdal.AllRegister()

        # Get the data directory from the environment
        st_data_dir = Sys.get_env_var('ST_DATA_DIR', None)

        # Call the main processing routine
        generate_unc(xml_filename=args.xml_filename,
                     st_data_dir=st_data_dir,
                     no_data_value=util.INTERMEDIATE_NO_DATA_VALUE,
                     num_threads=args.num_threads)

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

    logger.info('*** ST Generate UNC - Complete ***')

if __name__ == '__main__':
    main()