#!/opt/anaconda/bin/python

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

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 
ERR_NO_S1 = 50

# 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",
           ERR_NO_S1: "Could not find Sentinel-1 data with swath and polarisation provided"    
    }
 
    ciop.log(log_level, msg[exit_code])  

def process(s2_prds, s2_out_prd, ts = 1000, nc = 5, maxit = 1000, ct = 0.0001 ):

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

    OTB_app1 = otbApplication.Registry.CreateApplication("ConcatenateImages")    
    OTB_app1.SetParameterStringList("il", s2_prds)
    OTB_app1.SetParameterString('out', os.path.join(ciop.tmp_dir, 'concat.tif'))

    OTB_app1.ExecuteAndWriteOutput()

    # orchestrate the OTB applications
    OTB_app2 = otbApplication.Registry.CreateApplication('KMeansClassification')
    OTB_app2.SetParameterString('in', os.path.join(ciop.tmp_dir, 'concat.tif'))
    OTB_app2.SetParameterInt('ts', ts)
    OTB_app2.SetParameterInt('nc', nc)
    OTB_app2.SetParameterInt('maxit', maxit)
    OTB_app2.SetParameterString('out', os.path.join(ciop.tmp_dir, 'classif.tif'))

    OTB_app2.ExecuteAndWriteOutput()

    OTB_app3 = otbApplication.Registry.CreateApplication("ColorMapping")
    OTB_app3.SetParameterString("in", 'classif.tif')
    OTB_app3.SetParameterString("method","custom")
    OTB_app3.SetParameterString("method.custom.lut", '/application/otb/etc/classification.lut')
    OTB_app3.SetParameterString("out", s2_out_prd)
    OTB_app3.SetParameterOutputImagePixelType("out", 0)
    OTB_app3.ExecuteAndWriteOutput()

    os.remove(os.path.join(ciop.tmp_dir, 'concat.tif'))
    os.remove(os.path.join(ciop.tmp_dir, 'classif.tif'))
    

def main():

    os.chdir(ciop.tmp_dir)

    ts = int(ciop.getparam('ts'))
    nc = int(ciop.getparam('nc'))
    maxit = int(ciop.getparam('maxit'))
    ct = float(ciop.getparam('ct'))

    # 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)

      search_expression = os.path.join(ciop.tmp_dir, str(search[0]['identifier']), str(search[0]['identifier']) + '.SAFE', 'S2*.xml')

      ciop.log('INFO', 'path: ' + search_expression)

      #s1_prd = '' 
      try:
        s2_xml = glob.glob(search_expression)[0]  
      except:
        sys.exit(ERR_NO_S2_METADATA)

      os.chdir(os.path.join(ciop.tmp_dir, str(search[0]['identifier']), str(search[0]['identifier']) + '.SAFE', 'GRANULE'))

      granules = [ name for name in os.listdir(".") if os.path.isdir(name)]

      os.chdir(ciop.tmp_dir)
  
      for granule in granules:
        ciop.log('INFO', 'Processing granule ' + granule)

        b02 = os.path.join(ciop.tmp_dir, str(search[0]['identifier']), str(search[0]['identifier']) + '.SAFE', 'GRANULE', granule, 'IMG_DATA', granule[:-7] + '_B02.jp2')
        b03 = os.path.join(ciop.tmp_dir, str(search[0]['identifier']), str(search[0]['identifier']) + '.SAFE', 'GRANULE', granule, 'IMG_DATA', granule[:-7] + '_B03.jp2')
        b04 = os.path.join(ciop.tmp_dir, str(search[0]['identifier']), str(search[0]['identifier']) + '.SAFE', 'GRANULE', granule, 'IMG_DATA', granule[:-7] + '_B04.jp2')
        b08 = os.path.join(ciop.tmp_dir, str(search[0]['identifier']), str(search[0]['identifier']) + '.SAFE', 'GRANULE', granule, 'IMG_DATA', granule[:-7] + '_B08.jp2')
  
        s2_prds = [ b02, b03, b04, b08 ]      

        s2_result_prd = os.path.join(ciop.tmp_dir, granule[:-3] + '_kmeans.png')

        ciop.log('INFO', 'The output file is: ' + s2_result_prd)
        process(s2_prds, s2_result_prd, ts, nc, maxit, ct)
 
        assert(os.path.isfile(s2_result_prd)), sys.exit(ERR_NO_OUTPUT)
       
        # publish
        ciop.log('INFO', 'Publishing ' + s2_result_prd ) 
        ciop.publish(s2_result_prd, metalink=True)      
   
        # clean-up
        os.remove(s2_result_prd)
       
      # clean-up 
      shutil.rmtree(os.path.join(ciop.tmp_dir, str(search[0]['identifier'])))
try:
  main()
except SystemExit as e:
  if e.args[0]:
    clean_exit(e.args[0])
  raise
else:
  atexit.register(clean_exit, 0)


