#!/opt/anaconda/bin/python

import sys
import os
import lxml.etree as et
import shutil
import tarfile
import atexit

sys.path.append('/opt/anaconda/bin/')
import cioppy
ciop = cioppy.Cioppy()

sys.path.append(os.environ['_CIOP_APPLICATION_PATH'] + '/otbhelpers')
import otbhelpers as oh

sys.path.append('/opt/OTB/lib/python')
sys.path.append('/opt/OTB/lib/libfftw3.so.3')
os.environ['OTB_APPLICATION_PATH'] = '/opt/OTB/lib/otb/applications'
os.environ['LD_LIBRARY_PATH'] = '/opt/OTB/lib'
os.environ['ITK_AUTOLOAD_PATH'] = '/opt/OTB/lib/otb/applications'
import otbApplication

# define the exit codes
SUCCESS = 0
ERR_RESOLUTION = 10
ERR_STAGEIN = 20
ERR_EXTRACT = 30
ERR_NO_OUTPUT = 40 

# 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 Landsat-8 product enclosure',
           ERR_STAGEIN: 'Could not stage-in Landsat-8 product',
           ERR_EXTRACT: 'Failed to extract Landsat-8 product',
           ERR_NO_OUTPUT: "OTB failed to produce output"
    }
 
    ciop.log(log_level, msg[exit_code])  

def main():

    os.chdir(ciop.tmp_dir)

    paramfile1 = os.path.join(ciop.tmp_dir, 'DNtoTOA_1.xml' )
    paramfile2 = os.path.join(ciop.tmp_dir, 'DNtoTOA_2.xml' )
    paramfile3 = os.path.join(ciop.tmp_dir, 'NDVI.xml' )

    otbBandMath = os.environ['_CIOP_APPLICATION_PATH'] + '/otbhelpers/BandMath.xml'

    for paramfile in [ paramfile1, paramfile2, paramfile3 ]:
      shutil.copyfile(otbBandMath, paramfile) 

    # remove the XML block of parameter 'il' to use the method AddImageToParameterInputImageList
    tree = et.parse(paramfile3)
    root = tree.getroot()

    for exp in root.xpath('/OTB/application/parameter/key[text() = "il"]/..'):
      exp.getparent().remove(exp)
    
    et.ElementTree(root).write(paramfile3, pretty_print=True)

    # update the XML files
    gain = '2.0000E-05'
    offset = '-0.100000'
    
    for paramfile in [ paramfile1, paramfile2 ]:
      oh.set_OTB_param('exp', gain + ' * im1b1 + ' + offset, paramfile)

    oh.set_OTB_param('exp', "im1b1 >= 0 && im1b1 <= 1 && im2b1 >= 0 && im2b1 <= 1 ? ( im1b1 - im2b1 ) / ( im1b1 + im2b1 ) : 0 ", paramfile3)

    # vegetation index 
    index = ciop.getparam('vegetationindex') 
    # log the value, it helps debugging. 
    # the log entry is available in the process stderr 
    ciop.log('DEBUG', 'The index used is: ' + index)
    
    # Loops over all the inputs
    for inputfile in sys.stdin:
      # report activity in log
      ciop.log('INFO', 'The input file is: ' + inputfile)

      search = ciop.search(end_point = inputfile, params = [], output_fields='enclosure,identifier', model='GeoTime')
      assert(search), sys.exit(ERR_RESOLUTION)

      ciop.log('INFO', 'Retrieve %s from %s' % (search[0]['identifier'], search[0]['enclosure']))
      retrieved = ciop.copy(search[0]['enclosure'], ciop.tmp_dir)
      assert(retrieved), sys.exit(ERR_STAGEIN)  
 
      ciop.log('INFO', 'Extract %s' % (search[0]['identifier'])) 
      tar = tarfile.open(retrieved)
      tar.extractall()
      tar.close()

      try:
       band3 = glob.glob(ciop.tmp_dir + '/*_B3.TIF')[0]
       band4 = glob.glob(ciop.tmp_dir + '/*_B4.TIF')[0]
       band5 = glob.glob(ciop.tmp_dir + '/*_B5.TIF')[0]
       band6 = glob.glob(ciop.tmp_dir + '/*_B6.TIF')[0]
       band7 = glob.glob(ciop.tmp_dir + '/*_B7.TIF')[0]
      except:
       sys.exit(ERR_EXTRACT)
    
      assert(os.path.isfile(band3)), sys.exit(ERR_EXTRACT)
    
      ciop.log('INFO', 'Retrieved ' + os.path.basename(retrieved))

      if index == 'NDVI':
        oh.set_OTB_params('il', [ band5 ], paramfile1)
        oh.set_OTB_params('il', [ band4 ], paramfile2)

      if index == 'LSWI':
        oh.set_OTB_params('il', [ band5 ], paramfile1)
        oh.set_OTB_params('il', [ band6 ], paramfile2)

      if index == 'NBR':
        oh.set_OTB_params('il', [ band5 ], paramfile1)
        oh.set_OTB_params('il', [ band7 ], paramfile2)

      if index == 'MNDWI':
        oh.set_OTB_params('il', [ band3 ], paramfile1)
        oh.set_OTB_params('il', [ band6 ], paramfile2)  

      outputname = os.path.join(ciop.tmp_dir, search[0]['identifier'] + '_' + index + '.tif' )
      oh.set_OTB_param('out', outputname, paramfile3)

      ciop.log('INFO', 'Process in-memory OTB applications')

      # orchestrate the OTB applications
      OTB_app1 = otbApplication.Registry.CreateApplication(oh.get_OTB_appname(paramfile1))
      OTB_app1.SetParameterString("inxml", paramfile1)

      OTB_app1.Execute()

      OTB_app2 = otbApplication.Registry.CreateApplication(oh.get_OTB_appname(paramfile2))
      OTB_app2.SetParameterString("inxml", paramfile2)

      OTB_app2.Execute()

      OTB_app3 = otbApplication.Registry.CreateApplication(oh.get_OTB_appname(paramfile3))
      OTB_app3.SetParameterString("inxml", paramfile3)
      OTB_app3.AddImageToParameterInputImageList("il", OTB_app1.GetParameterOutputImage("out"))
      OTB_app3.AddImageToParameterInputImageList("il", OTB_app2.GetParameterOutputImage("out"))

      OTB_app3.ExecuteAndWriteOutput()
  
      assert(os.path.isfile(outputname)), sys.exit(ERR_NO_OUTPUT)
       
      # publish
      ciop.log('INFO', 'Publishing ' + outputname ) 
      ciop.publish(outputname, metalink=True)      

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

