#!/opt/anaconda/bin/python

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

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

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'
os.environ['GDAL_HOME'] = glob.glob('/opt/anaconda/pkgs/gdal-2.?.?-py27_?')[0]
os.environ['GDAL_DATA'] = '/opt/OTB/share/gdal'
os.environ['GEOTIFF_CSV'] = '/opt/OTB/share/epsg_csv'
os.environ['LC_NUMERIC'] = 'C'

import otbApplication

# define the exit codes
SUCCESS = 0
ERR_RESOLUTION = 10
ERR_STAGEIN = 20
ERR_EXTRACT = 30
ERR_NO_OUTPUT = 40 
ERR_NO_KOMPSAT_META = 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 KOMPSAT-3 product enclosure',
           ERR_STAGEIN: 'Could not stage-in KOMPSAT-3 product',
           ERR_EXTRACT: 'Failed to extract KOMPSAT-3 product',
           ERR_NO_OUTPUT: "OTB failed to produce output",
           ERR_NO_KOMPSAT_META: "Could not find KOMPSAT-3 metadata file"    
    }
 
    ciop.log(log_level, msg[exit_code])  

def process(red, green, blue, nir, pan, aux, k3_out_prd, level = 'toa'):

    ciop.log('INFO', 'Process OTB applications')

    k3_concat = os.path.join(ciop.tmp_dir, '_concat.tif')
    k3_pansharp = os.path.join(ciop.tmp_dir, '_pansharp.tif')

    k3_prds = [ red, green, blue ]

    type_data = otbApplication.ImagePixelType_float
    type_visu = otbApplication.ImagePixelType_uint8

    tree = et.parse(aux)
    root = tree.getroot()

    azimuths = list()
    elevations = list()

    for azimuth in root.iter('Azimuth'):
      azimuths.append(float(azimuth.text))

    for elevation in root.iter('Elevation'):
      elevations.append(float(elevation.text))

    azi = np.mean(azimuths)
    ele = np.mean(elevations)

    year = int(os.path.basename(red)[3:7])
    month = int(os.path.basename(red)[7:9])
    day = int(os.path.basename(red)[9:11])
    hour = int(os.path.basename(red)[11:13])
    minute = int(os.path.basename(red)[13:15])

    gainbias = [ '/application/otb/etc/k3_gainbias_red.txt', 
                 '/application/otb/etc/k3_gainbias_green.txt', 
                 '/application/otb/etc/k3_gainbias_blue.txt' ]
    
    solarillumination = [ '/application/otb/etc/k3_solarillumination_red.txt', 
                          '/application/otb/etc/k3_solarillumination_green.txt', 
                          '/application/otb/etc/k3_solarillumination_blue.txt' ]

    tmp_oc = [ os.path.join(ciop.tmp_dir, '_oc_red.tif'),
               os.path.join(ciop.tmp_dir, '_oc_green.tif'),
               os.path.join(ciop.tmp_dir, '_oc_blue.tif')]

    for index, item in enumerate(k3_prds, start = 0):
      OTB_app1 = otbApplication.Registry.CreateApplication("OpticalCalibration")
      OTB_app1.SetParameterString("in", k3_prds[index])
      OTB_app1.SetParameterString("level", level)
      OTB_app1.SetParameterInt('ram', 4096)
      OTB_app1.SetParameterString('acqui.gainbias', gainbias[index])
      OTB_app1.SetParameterString('acqui.solarilluminations', solarillumination[index])
      OTB_app1.SetParameterFloat('acqui.sun.elev', ele)
      OTB_app1.SetParameterFloat('acqui.sun.azim', azi)
      OTB_app1.SetParameterInt('acqui.minute', minute)
      OTB_app1.SetParameterInt('acqui.hour', hour)
      OTB_app1.SetParameterInt('acqui.day', day)
      OTB_app1.SetParameterInt('acqui.month', month)
      OTB_app1.SetParameterInt('acqui.year', year)
      OTB_app1.SetParameterString('atmo.aerosol', 'noaersol')
      OTB_app1.SetParameterString('out', tmp_oc[index])
      OTB_app1.SetParameterString('outxml', tmp_oc[index] + '.xml')
      OTB_app1.SetParameterString("milli", '1')
      OTB_app1.DisableParameter('clamp')
      OTB_app1.ExecuteAndWriteOutput()

      ciop.publish(tmp_oc[index] + '.xml', metalink=True)
 
    ciop.log('INFO', 'Concatenate optical calibrated bands')
    OTB_app2 = otbApplication.Registry.CreateApplication("ConcatenateImages")
    OTB_app2.SetParameterStringList("il", tmp_oc)
    OTB_app2.SetParameterOutputImagePixelType("out", type_data)
    OTB_app2.SetParameterInt('ram', 4096)
    OTB_app2.SetParameterString('out', k3_concat)

    OTB_app2.ExecuteAndWriteOutput()

    ciop.log('INFO', 'Pansherpening')
    OTB_app3 = otbApplication.Registry.CreateApplication("BundleToPerfectSensor")
    OTB_app3.SetParameterString("inp", pan)
    OTB_app3.SetParameterString("inxs", k3_concat)
    OTB_app3.SetParameterInt("ram", 4096)
    OTB_app3.SetParameterString("out", k3_pansharp)
    OTB_app3.SetParameterOutputImagePixelType("out", type_data)

    OTB_app3.ExecuteAndWriteOutput()

    ciop.log('INFO', 'Cconvertion to uint8 data type')
    OTB_app4 = otbApplication.Registry.CreateApplication("Convert")
    OTB_app4.SetParameterString("in", k3_pansharp)
    OTB_app4.SetParameterString('type', 'linear')
    OTB_app4.SetParameterOutputImagePixelType("out", type_visu)
    #OTB_app4.SetParameterString('channels', 'rgb')
    OTB_app4.SetParameterString('out', k3_out_prd)
    OTB_app4.ExecuteAndWriteOutput() 

    ciop.log('INFO', 'Clean-up')
    for tmp_prd in tmp_oc:
    	os.remove(os.path.join(ciop.tmp_dir, tmp_prd ))
    
    os.remove(k3_concat)
    os.remove(k3_pansharp)
  
def main():

    os.chdir(ciop.tmp_dir)

    level = ciop.getparam('level')

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

      k3_prd_r = [os.path.join(dirpath, f)
      for dirpath, dirnames, files in os.walk(retrieved)
      for f in fnmatch.filter(files, 'K3*L1G_R.tif')][0]

      k3_prd_g = [os.path.join(dirpath, f)
        for dirpath, dirnames, files in os.walk(retrieved)
        for f in fnmatch.filter(files, 'K3*L1G_G.tif')][0]

      k3_prd_b = [os.path.join(dirpath, f)
        for dirpath, dirnames, files in os.walk(retrieved)
        for f in fnmatch.filter(files, 'K3*L1G_B.tif')][0]

      k3_prd_n = [os.path.join(dirpath, f)
        for dirpath, dirnames, files in os.walk(retrieved)
        for f in fnmatch.filter(files, 'K3*L1G_N.tif')][0]

      k3_prd_p = [os.path.join(dirpath, f)
        for dirpath, dirnames, files in os.walk(retrieved)
        for f in fnmatch.filter(files, 'K3*L1G_P.tif')][0]

      k3_aux = [os.path.join(dirpath, f)
        for dirpath, dirnames, files in os.walk(retrieved)
        for f in fnmatch.filter(files, 'K3*L1G_Aux.xml')][0]

      k3_result_prd = k3_prd_r[:-10] + '_opt_cal.tif'

      try:
        ciop.log('INFO', 'KOMPSAT-3 metadata file: ' + k3_aux)
      except:
        sys.exit(ERR_NO_KOMPSAT_META)
 
      ciop.log('INFO', 'The output file is: ' + k3_result_prd)
      process(k3_prd_r, k3_prd_g, k3_prd_b, k3_prd_n, k3_prd_p, k3_aux, k3_result_prd, level)
  
      assert(os.path.isfile(k3_result_prd)), sys.exit(ERR_NO_OUTPUT)
       
      # publish
      ciop.log('INFO', 'Publishing ' + k3_result_prd ) 
      ciop.publish(k3_result_prd, metalink=True)      
   
      # clean-up
      os.remove(k3_result_prd)
       
      # clean-up 
      shutil.rmtree(retrieved)
try:
  main()
except SystemExit as e:
  if e.args[0]:
    clean_exit(e.args[0])
  raise
else:
  atexit.register(clean_exit, 0)


