#!/opt/anaconda/bin/python

import sys
reload(sys)
sys.setdefaultencoding('utf8')
import os
import io
import shutil
import atexit
from nbconvert.preprocessors import ExecutePreprocessor, CellExecutionError
import nbformat as nbf
import uuid
import ast

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

# define the exit codes
SUCCESS = 0
ERR_NB_RUNTIME=10

references = []

# 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_NB_RUNTIME: 'Failed to run notebook'
    }
 
    ciop.log(log_level, msg[exit_code])  

def parametrize():
    global nb
        
    for index, cell in enumerate(nb['cells']):
       
        if str(cell['cell_type']) == 'code': 
        
            try:
                root_ast = ast.parse(str(cell['source']))
                names = list({node.id for node in ast.walk(root_ast) if isinstance(node, ast.Name)})
                
                if len(names) == 1:
                    
                    if names[0] == 'data_path':
                        ciop.log('INFO', 'cell %s updated with \'data_path\' value %s' % (index, tmp_dir))
                        cell['source'] = 'data_path = \'%s\'' % tmp_dir  
                
                    if names[0] == 'input_identifiers':
                        ciop.log('INFO', 'cell %s updated with \'input_identifiers\' value %s' % (index, identifiers))
                        cell['source'] = 'input_identifiers = %s' % identifiers  
                
                    if names[0] == 'input_references':
                        ciop.log('INFO', 'cell %s updated with \'input_references\' value %s' % (index, references))
                        cell['source'] = 'input_references = %s' % references  
                
                if len(names) != 2:
                    continue
                
                if names[0] == 'dict' or names[1] == 'dict':
                    
                    # deal with the alphabetical order
                    if names[1] == 'dict': 
                        names[1] = names[0]
                        names[0] = 'dict'
                
                    exec(str(cell['source'])) in globals(), locals()
                
                    if names[0] == 'dict' and 'title' in eval(names[1]).keys() and 'abstract' in eval(names[1]).keys() and 'id' in eval(names[1]).keys() and 'value' in eval(names[1]).keys():
                                      
                        eval(names[1])['value'] = ciop.getparam(eval(names[1])['id'])
    
                        new_source = 'dict(['

                        for i, keys in enumerate(eval(names[1])):
                            if i == 0: 
                                new_source = new_source + '( "%s", "%s")' % (keys, eval(names[1])[keys]) 
                            else:
                                new_source = new_source + ',( "%s", "%s")' % (keys, eval(names[1])[keys]) 
        
                        new_source = new_source + '])'
    
                        cell['source'] = '%s = %s' % (names[1], new_source)
         
                        ciop.log('INFO', 'cell %s %s updated' % (index, names[1]))
                    
            except SyntaxError:
                continue   



def execute(nb_source, nb_target, kernel = 'python2'):
    
    global nb
   
    nb = nbf.read(nb_source, 4)
    
    ciop.log('INFO', 'Execute notebook')
    
    parametrize()
    
    # execute the notebook
    ep = ExecutePreprocessor(timeout=None, kernel_name=kernel)

    try:
      out = ep.preprocess(nb, {'metadata': {'path': './'}})
    except CellExecutionError:
      out = None
      ciop.log('ERROR', 'Error executing the notebook "%s".' % nb_source)

      with io.open(nb_target, 'wb') as file:
        file.write(nbf.writes(nb))
        file.close()
      
      ciop.publish(nb_target, metalink=True)
      raise
    finally:
      ciop.log('INFO', 'Write notebook')  
      with io.open(nb_target, 'wb') as file:
        file.write(nbf.writes(nb))
        file.close() 

def publish():
   
    # publish
    ciop.log('INFO', 'Publishing ') 
    ciop.publish(runtime, metalink=True, recursive=True)

def clean_up():
           
    # clean-up 
    shutil.rmtree(runtime)

def main():

    # create the folder for the data stage-in
    global runtime
    runtime = os.path.join(ciop.tmp_dir, str(uuid.uuid4()))    

    ciop.log('DEBUG', 'The runtime folder is %s' % (runtime))
    os.makedirs(runtime)
    os.chdir(runtime) 

    # Loops over all the inputs
    for reference in sys.stdin:
      references.append(reference.rstrip())

    
    # execute the notebook
    nb_source = os.path.join('/application', 'notebook', 'libexec', 'input.ipynb')
    nb_target = os.path.join(runtime, 'result.ipynb')  
    execute(nb_source, nb_target, 'python2')  
    
    # publish
    publish()   
   
    # clean-up 
    clean_up()

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


