#!/opt/anaconda/envs/env_ewf_satcen_03_01_01/bin/python

# This node reads dim files coherence and sigma0 for sub-swaths & polarisation/s 
# Merge subswaths together & speckle filters, then merge sigma0 bands and coherence 
# post-processing  gdal calculations



import os
import sys
import string
import atexit
import tarfile
import cioppy 
ciop = cioppy.Cioppy()

from snappy import jpy
from snappy import ProductIO
sys.path.append('/'.join([os.environ['_CIOP_APPLICATION_PATH'], 'util']))
sys.path.append('../util')
from py_snap_helpers import *
from util import *
sys.path.append('/opt/anaconda/envs/env_ewf_satcen_03_01_01/snap/.snap/snap-python')    
# define the exit codes
SUCCESS = 0
ERR_RESOLUTION = 10
ERR_STAGEIN = 20
ERR_NO_OUTPUT = 30

# add a trap to exit gracefully
def clean_exit(exit_code):
    log_level = 'INFO'
    if exit_code != SUCCESS:
        log_level = 'ERROR'  
   
    msg = {SUCCESS: 'Processing successfully concluded',
           ERR_RESOLUTION: 'Could not resolve Sentinel-1 product enclosure',
           ERR_STAGEIN: 'Could not stage-in Sentinel-1 product', 
           ERR_NO_OUTPUT: "Missing output"
    }
 
    ciop.log(log_level, msg[exit_code])  

def main():
        
    ciop = cioppy.Cioppy()
    
    os.chdir(ciop.tmp_dir)
    ciop.log('DEBUG', 'ciop path is :{}'.format(ciop.tmp_dir))
    
    #aoi = ciop.getparam('aoi')
    coh_input_splits = []
    sigma0_input_splits = []
    
    #### Read all input zip files consisting of two types dim files: coherence and sigma0
    input_extract_dir = os.path.join(ciop.tmp_dir,'input_dir')
    if not os.path.exists(input_extract_dir):
        os.makedirs(input_extract_dir)
 

    for input in sys.stdin:
        ciop.log('INFO', 'Processing file {}'.format(input.rstrip()))
        #Only zip data files extracted 
        if input.rstrip().split('.')[-1] == 'gz':
            ciop.copy(input.rstrip(),input_extract_dir,extract=False)
            ciop.log('DEBUG','zipped input {} is staged-in'.format(os.path.basename(input.rstrip())))
            with tarfile.open(os.path.join(input_extract_dir,os.path.basename(input.rstrip())), "r:gz") as input_tarfile:
                input_tarfile.extractall(path=input_extract_dir)
            os.remove(os.path.join(input_extract_dir,os.path.basename(input.rstrip())))
        
        
    ## Discriminate input files by name: Choose dim files to be inserted in SNAP graph
    for file in os.listdir(input_extract_dir):
        ciop.log('DEBUG', 'Processing file {}'.format(file))
        #if file[-7:-4] == 'coh' in and file.split('.')[-1]=='dim':
        if 'coh' in file and '.dim' in file:
            coh_input_splits.append(os.path.join(input_extract_dir,file))
            
        #elif file[-10:-4] == 'sigma0' and file.split('.')[-1]=='dim':
        elif 'sigma0' in file and '.dim' in file:
            sigma0_input_splits.append(os.path.join(input_extract_dir,file))
            output_name = file[:-11]
            
    ciop.log('INFO', 'Total number of {} sub-swath & Polarisation input slices.'.format(len(coh_input_splits)))
        
    #####Process Inputs
    graph = GraphProcessor('/opt/anaconda/envs/env_ewf_satcen_03_01_01/snap/bin/gpt')
    coh_read_nodes = []
    sigma0_read_nodes = []
    #### Process coherence input splits
    operator = 'Read'
    
    for i in range(len(coh_input_splits)):
        
        parameters = get_operator_default_parameters(operator)
        node_id = 'Read-COH-{}'.format(i)
        source_node_id = ''
        parameters['file'] = coh_input_splits[i]
        graph.add_node(node_id,
                         operator, 
                         parameters,
                         source_node_id)
        
        coh_read_nodes.append(node_id)
        source_node_id_coh = node_id
        
    for i in range(len(sigma0_input_splits)):
        
        parameters = get_operator_default_parameters(operator)
        node_id = 'Read-sigma0-{}'.format(i)
        source_node_id = ''
        parameters['file'] = sigma0_input_splits[i]
        graph.add_node(node_id,
                         operator, 
                         parameters,
                         source_node_id)
        
        sigma0_read_nodes.append(node_id)
        source_node_id_sigma0 = node_id
        
    #### if one sub-swath skip TOPSAR-Merge   
    if  len(coh_input_splits) > 1:
        source_nodes_id = coh_read_nodes
        operator = 'TOPSAR-Merge'
        node_id = 'TOPSAR-Merge_coh_inputs'
        parameters = get_operator_default_parameters(operator)
        #parameters['selectedPolarisations'] = polarization
        graph.add_node(node_id,
                         operator, 
                         parameters,
                         source_nodes_id)
        source_node_id_coh = node_id
       
        
    if  len(sigma0_input_splits) > 1:
        source_nodes_id = sigma0_read_nodes
        operator = 'TOPSAR-Merge'
        node_id = 'TOPSAR-Merge_sigma0_inputs'
        parameters = get_operator_default_parameters(operator)
        #parameters['selectedPolarisations'] = polarization  
        graph.add_node(node_id,
                         operator, 
                         parameters,
                         source_nodes_id)
        source_node_id_sigma0 = node_id
            
            
    ####Speckle filtering sigma0 merged inputs
    operator = 'Speckle-Filter'
    parameters = get_operator_default_parameters(operator)
    parameters['filter'] = 'Lee Sigma'
    parameters['filterSizeX'] = '3'
    parameters['filterSizeY'] = '3' 
    parameters['dampingFactor'] = '2' 
    parameters['estimateENL'] = 'true'
    parameters['enl'] = '1.0'
    parameters['numLooksStr'] = '1'
    parameters['windowSize'] = '7x7'
    parameters['targetWindowSizeStr'] = '3x3'
    parameters['sigmaStr'] = '0.9'
    parameters['anSize'] = '50'
        
    node_id = 'Speckle-Filter_sigma0_inputs'         
    graph.add_node(node_id,
                         operator, 
                         parameters,
                         source_node_id_sigma0)
    source_node_id_sigma0 = node_id
    
    ####Band Merge
    operator = 'BandMerge'
    parameters = get_operator_default_parameters(operator)
    parameters['geographicError'] = '1.0E-5'
        
    node_id = 'BandMerge'         
    graph.add_node(node_id,
                         operator, 
                         parameters,
                         [source_node_id_sigma0,source_node_id_coh])  
    source_node_id = node_id
    
    ####
    operator = 'Terrain-Correction'
    parameters = get_operator_default_parameters(operator)
    parameters['demName'] = ciop.getparam('demtype')#'SRTM 1Sec HGT'
    parameters['externalDEMNoDataValue'] = '0.0'
    parameters['externalDEMApplyEGM'] = 'true' 
    parameters['demResamplingMethod'] = 'BILINEAR_INTERPOLATION' 
    parameters['imgResamplingMethod'] = 'BILINEAR_INTERPOLATION'
    parameters['pixelSpacingInMeter'] = ciop.getparam('pixelSpacingInMeter')
    parameters['pixelSpacingInDegree'] = '0'
    parameters['mapProjection'] = ciop.getparam('mapProjection')
    parameters['alignToStandardGrid'] = 'false'
    parameters['standardGridOriginX'] = '0'
    parameters['standardGridOriginY'] = '0'
    parameters['nodataValueAtSea'] = 'true'
    parameters['saveDEM'] = 'false'
    parameters['saveLatLon'] = 'false'
    parameters['saveIncidenceAngleFromEllipsoid'] = 'false'
    parameters['saveLocalIncidenceAngle'] = 'false'
    parameters['saveProjectedLocalIncidenceAngle'] = 'false'
    parameters['saveSelectedSourceBand'] = 'true'
    parameters['outputComplex'] = 'false'
    parameters['applyRadiometricNormalization'] = 'false'
    parameters['saveSigmaNought'] = 'false'
    parameters['saveGammaNought'] = 'false'
    parameters['saveBetaNought'] = 'false'
    parameters['incidenceAngleForSigma0'] = 'Use projected local incidence angle from DEM'
    parameters['incidenceAngleForGamma0'] = 'Use projected local incidence angle from DEM'
    parameters['auxFile'] = 'Latest Auxiliary File'
    
    node_id = 'Terrain-Correction'         
    graph.add_node(node_id,
                         operator, 
                         parameters,
                         source_node_id)
    source_node_id = node_id
    
    #### Create an output_dir to place the output
#    local_output_dir = os.path.join(ciop.tmp_dir,'local_output_dir')
#    if not os.path.exists(local_output_dir):
#        os.makedirs(local_output_dir)

#----------------------------------------------------------------
#    ####
    operator = 'Subset'
    source_node_id = node_id
    node_id = 'Subset'

    parameters = get_operator_default_parameters(operator)
    parameters['geoRegion'] = ciop.getparam('aoi')
    parameters['copyMetadata'] = 'true'


    graph.add_node(node_id,
                   operator,
                   parameters,
                   source_node_id)
    
    source_node_id = node_id


    ####    
    operator = 'Write'
    parameters = get_operator_default_parameters(operator)
    #parameters['formatName'] = 'BEAM-DIMAP'
    parameters['formatName'] = 'GeoTIFF-BigTIFF'
    parameters['file'] = os.path.join(ciop.tmp_dir, 'target.tif')    
    node_id = 'Write'         
    graph.add_node(node_id,
                   operator, 
                   parameters,
                   source_node_id)
    
    
    ####Run Graph
    graph.run() 
    
    ciop.log('INFO', 'SNAP Processing finished successfully.')

    # For test publish also raw output data
    #ciop.publish(os.path.join(ciop.tmp_dir, 'target.tif'), metalink=True)
    
    ####Gdal post-processing to be added

    master_info = '_'.join(os.path.basename(sigma0_input_splits[0]).split('_')[:6])
    slave_info = '_'.join(os.path.basename(sigma0_input_splits[0]).split('_')[6:12])

    
    master_date = os.path.basename(sigma0_input_splits[0]).split('_')[5]
    slave_date = os.path.basename(sigma0_input_splits[0]).split('_')[11]
    
    ##Sigma0_{master_info}; Sigma0_{slave_info}
    Sigma0_master = 'Sigma0_{}.tif'.format(master_info)
    Sigma0_slave = 'Sigma0_{}.tif'.format(slave_info)
    
    ciop.log('INFO','Creating Sigma0 master COG')
    cog(os.path.join(ciop.tmp_dir,'target.tif'), os.path.join(ciop.tmp_dir,Sigma0_master), band=1)
    ciop.log('INFO','Creating Sigma0 slave COG')
    cog(os.path.join(ciop.tmp_dir,'target.tif'), os.path.join(ciop.tmp_dir,Sigma0_slave), band=2)
    
    
    ##Coherence_{master_info}_{slave_info}
    Coh = 'Coherence_{}_{}.tif'.format(master_info, slave_info)
    ciop.log('INFO','Creating Coherence COG')
    cog(os.path.join(ciop.tmp_dir,'target.tif'),os.path.join(ciop.tmp_dir,Coh), band=3)
    
    
    ##ACD RGB file, name: ACD_{master_date}_{slave_date}, R: Sigma0_{master_info}, G: Sigma0_{slave_info}, B: Sigma0_{slave_info}
    ACD = 'ACD_{}_{}.tif'.format(master_date,slave_date)
    ciop.log('INFO','Creating ACD COG')
    COG_merge(os.path.join(ciop.tmp_dir,Sigma0_master),os.path.join(ciop.tmp_dir,Sigma0_slave),os.path.join(ciop.tmp_dir,Sigma0_slave),os.path.join(ciop.tmp_dir,ACD))
    
    
    ##CCD RGB file, name: CCD_{master_date}_{slave_date},  R: Sigma0_{master_info}, G: Sigma0_{slave_info}, B: Coherence_{master_info}_{slave_info}
    CCD = 'CCD_{}_{}.tif'.format(master_date,slave_date)
    ciop.log('INFO','Creating CCD COG')
    COG_merge(os.path.join(ciop.tmp_dir,Sigma0_master),os.path.join(ciop.tmp_dir,Sigma0_slave),os.path.join(ciop.tmp_dir,Coh),os.path.join(ciop.tmp_dir,CCD))
    
    
    #### master/slave date conversion in ISO format
    date_m = master_date.split('T')[0]
    time_m = master_date.split('T')[1]

    master_date = '{}-{}-{}T{}:{}:{}Z'.format(date_m[:4],date_m[4:6],date_m[6:8],time_m[:2],time_m[2:4],time_m[4:])
    date_s= slave_date.split('T')[0]
    time_s = slave_date.split('T')[1]
    slave_date = '{}-{}-{}T{}:{}:{}Z'.format(date_s[:4],date_s[4:6],date_s[6:8],time_s[:2],time_s[2:4],time_s[4:])

    ####Publish Results
    for index, output_file in enumerate([ Sigma0_master,
                                          Sigma0_slave,
                                          Coh,
                                          ACD,
                                          CCD]):
                                            
        if output_file == Sigma0_master:
            title = 'Sigma0 - Backscatter coefficient for master {0}'.format(master_info)
            date = 'date={0}/{1}'.format(master_date,
                                         master_date)
        elif output_file == Sigma0_slave:
            title = 'Sigma0 - Backscatter coefficient for slave {0}'.format(slave_info)
            date = '{0}/'.format(slave_date,
                                 slave_date)

        elif output_file == Coh:
            title = 'Sentinel1 coherence between master {0} and slave {1}'.format(master_info,
                                                                                  slave_info)
            date = '{0}/{1}'.format(master_date,
                                     slave_date)
        elif output_file == CCD:
            title = 'RGB composite of the coherence and backscatter coefficient for master/slave {}/{}'.format(master_info,
                                                                                                               slave_info)
            date = '{0}/{1}'.format(master_info,
                                    slave_date)
        elif output_file == ACD:
            title = 'RGB composite of the backscatter coefficient for master/slave {}/{}'.format(master_info,
                                                                                                 slave_info)
            date = '{0}/{1}'.format(master_info,
                                    slave_date)

        with open(os.path.join(ciop.tmp_dir, '{}.properties'.format(output_file)) , 'wb') as file:
            file.write(bytes('title={0}\n'.format(title),encoding = 'utf-8'))
            file.write(bytes('date={0}\n'.format(date),encoding = 'utf-8'))
            if (index > 1):
                file.write(bytes('geometry=',encoding = 'utf-8'))
                file.write(bytes(get_image_wkt(os.path.join(ciop.tmp_dir, output_file)),encoding = 'utf-8'))


        #create_metadata(output_tif) #No refrence to function found
        # Beta Version: result & stage-in removed *format to be controlled*

        ciop.publish(os.path.join(ciop.tmp_dir, output_file), metalink=True)
        ciop.publish(os.path.join(ciop.tmp_dir, '{}.properties'.format(output_file)), metalink=True)
        
        

try:
    main()
except SystemExit as e:
    if e.args[0]:
        clean_exit(e.args[0])
    raise
else:
    atexit.register(clean_exit, 0)
