#!/usr/bin/env python3
"""
Authors: Ioannis Paraskevakos, Aymen Alsaadi
License: MIT
Copyright: 2018-2019
"""

import radical.utils as ru
from iceberg.iceberg_parser import IcebergParser

from iceberg.executor import Seals
from iceberg.executor import Penguins
from iceberg.executor import Rivers

if __name__ == "__main__":

    parsed_values = IcebergParser().args()
    import os
    if parsed_values['general'].get('rmq_username', None):
        os.environ['RMQ_USERNAME'] = parsed_values['general']['rmq_username']
    elif 'RMQ_USERNAME' not in os.environ:
        raise RuntimeError('RMQ_USERNAME is not setup.')
    
    if parsed_values['general'].get('rmq_password', None):
        os.environ['RMQ_PASSWORD'] = parsed_values['general']['rmq_password']
    elif 'RMQ_PASSWORD' not in os.environ:
        raise RuntimeError('RMQ_PASSWORD is not setup.')

    if parsed_values['general'].get('rmq_endpoint', None):
        os.environ['RMQ_ENDPOINT'] = parsed_values['general']['rmq_endpoint']
    elif 'RMQ_ENDPOINT' not in os.environ:
        raise RuntimeError('RMQ_ENDPOINT is not setup.')
    
    if parsed_values['general'].get('rmq_port', None):
        os.environ['RMQ_PORT'] = parsed_values['general']['rmq_port']
    elif 'RMQ_PORT' not in os.environ:
        raise RuntimeError('RMQ_PORT is not setup.')
    
    if parsed_values['general'].get('radical_pilot_dburl', None):
        os.environ['RADICAL_PILOT_DBURL'] = parsed_values['general']['radical_pilot_dburl']
    elif 'RADICAL_PILOT_DBURL' not in os.environ:
        raise RuntimeError('RADICAL_PILOT_DBURL is not setup.')
    
    if parsed_values['analysis']['which'] == 'seals':
        if parsed_values['analysis'].get('ve_seals', None):
            os.environ['VE_SEALS'] = parsed_values['analysis']['ve_seals']
        elif 'VE_SEALS' not in os.environ:
            raise RuntimeError('VE_SEALS is not setup.')
        exec_obj = Seals(name=ru.generate_id('seals',
                         mode=ru.ID_PRIVATE),
                         resources={
                            'resource': parsed_values['general']['resource'],
                            'queue': parsed_values['general']['queue'],
                            'walltime': parsed_values['general']['walltime'],
                            'cpus': parsed_values['general']['cpus'],
                            'gpus': parsed_values['general']['gpus']
                            },
                         project=parsed_values['general']['project'],
                         input_path=parsed_values['general']['input_path'],
                         output_path=parsed_values['general']['output_path'],
                         bands=parsed_values['analysis']['bands'],
                         stride=parsed_values['analysis']['stride'],
                         patch_size=parsed_values['analysis']['patch_size'],
                         geotiff=parsed_values['analysis']['geotiff'],
                         model_arch=parsed_values['analysis']['model_architecture'],
                         hyperparam_set=parsed_values['analysis']['hyperparameter_set'],
                         model_name=parsed_values['analysis']['model_name'],
                         models_folder=parsed_values['analysis']['models_folder'])

    elif parsed_values['analysis']['which'] == 'penguins':
        if parsed_values['analysis'].get('ve_penguins', None):
            os.environ['VE_PENGUINS'] = parsed_values['analysis']['ve_penguins']
        elif 'VE_PENGUINS' not in os.environ:
            raise RuntimeError('VE_PENGUINS is not setup.')
        exec_obj = Penguins(name=ru.generate_id('penguins',
                         mode=ru.ID_PRIVATE),
                         resources={
                            'resource': parsed_values['general']['resource'],
                            'queue': parsed_values['general']['queue'],
                            'walltime': parsed_values['general']['walltime'],
                            'cpus': parsed_values['general']['cpus'],
                            'gpus': parsed_values['general']['gpus']
                            },
                         project=parsed_values['general']['project'],
                         input_path=parsed_values['general']['input_path'],
                         output_path=parsed_values['general']['output_path'],
                         model=parsed_values['analysis']['model'],
                         model_path=parsed_values['analysis']['model_path'], 
                         epoch = parsed_values['analysis']['epoch'])

    elif parsed_values['analysis']['which'] == 'rivers':
        if parsed_values['analysis'].get('ve_rivers', None):
            os.environ['VE_RIVERS'] = parsed_values['analysis']['ve_rivers']
        elif 'VE_RIVERS' not in os.environ:
            raise RuntimeError('VE_RIVERS is not setup.')
        exec_obj = Rivers(name=ru.generate_id('rivers',
                          mode=ru.ID_PRIVATE),
                          resources={
                            'resource': parsed_values['general']['resource'],
                            'queue': parsed_values['general']['queue'],
                            'walltime': parsed_values['general']['walltime'],
                            'cpus': parsed_values['general']['cpus'],
                            'gpus': parsed_values['general']['gpus']
                            },
                          project=parsed_values['general']['project'],
                          input_path=parsed_values['general']['input_path'],
                          output_path=parsed_values['general']['output_path'],
                          tile_size=parsed_values['analysis']['tile_size'],
                          step=parsed_values['analysis']['step'],
                          weights_path=parsed_values['analysis']['weights_path'])

    else:
        raise RuntimeError('Analysis %s not supported yet' %
                           parsed_values['analysis']['which'])

    exec_obj.run()
