#! /usr/bin/env python

import argparse
import json
import sys, os

try:
    import hpogrid
except:
    raise ImportError('Cannot import hpogrid module. Try source setupenv.sh first.')

kCommand_HPOConfig = 'hpo_config'
kCommand_GridConfig = 'grid_config'
kCommand_ModelConfig = 'model_config'
kCommand_SearchSpace = 'search_space'
kCommand_Project = 'project'
kCommand_PandaStatus = 'tasks'
kCommand_RunProject = 'run'
kCommand_RunProjectLocally = 'local_run'
kCommand_RunProjectInsideGrid = 'grid_run'
kCommand_RunProjectWithiDDS = 'idds_run'
kCommand_SubmitiDDSJob = 'submit_idds'
kCommand_SiteInfo = 'sites'
kCommand_Report = 'report'
kCommand_Generate = 'generate'

kCommandList = [kCommand_HPOConfig, kCommand_GridConfig, kCommand_ModelConfig,
                kCommand_SearchSpace, kCommand_Project, kCommand_PandaStatus,
                kCommand_SiteInfo, kCommand_RunProjectLocally,
                kCommand_RunProjectInsideGrid, kCommand_RunProjectWithiDDS,
                kCommand_SubmitiDDSJob,
                kCommand_RunProject, kCommand_Report, kCommand_Generate]


class HPOGridCLI(object):
    def __init__(self):
        parser = argparse.ArgumentParser(
            description='Tool for Hyperparameter Optimization on the Grid',
            usage='''hpogrid <command> [<args>]''')
        parser.add_argument('command', help='Subcommand to run', choices=kCommandList)
        # parse_args defaults to [1:] for args, but you need to
        # exclude the rest of the args too, or validation will fail        
        args = parser.parse_args(sys.argv[1:2])
        if not hasattr(self, args.command):
                    print('Unrecognized command')
                    parser.print_help()
                    exit(1)
        getattr(self, args.command)()

    def hpo_config(self):
        from hpogrid.configuration.hpo_configuration import HPOConfiguration
        HPOConfiguration()

    def grid_config(self):
        from hpogrid.configuration.grid_configuration import GridConfiguration
        GridConfiguration()

    def search_space(self):
        from hpogrid.configuration.search_space_configuration import SearchSpaceConfiguration
        SearchSpaceConfiguration()

    def model_config(self):
        from hpogrid.configuration.model_configuration import ModelConfiguration
        ModelConfiguration()

    def project(self):
        from hpogrid.configuration.project_configuration import ProjectConfiguration
        ProjectConfiguration()
    
    def tasks(self):
        from hpogrid.components.panda_task_manager import PandaTaskManager
        PandaTaskManager(True)

    def run(self):
        from hpogrid.components import GridHandler
        GridHandler.GridHandler()

    def local_run(self):
        from hpogrid.components import job_builder
        job_builder.local_job_parser()
        
    def grid_run(self):
        from hpogrid.components import job_builder
        job_builder.grid_job_parser() 

    def idds_run(self):
        from hpogrid.components import job_builder
        job_builder.grid_job_parser(idds_job=True) 
    
    def submit_idds(self):
        from hpogrid.idds_interface.idds_handler import iDDSHandler
        iDDSHandler()
        
    def sites(self):
        from hpogrid.utils.grid_site_info import GridSiteInfo
        GridSiteInfo()

    def report(self):
        from hpogrid.components.reporter import HPOTaskHandle
        HPOTaskHandle()
        
    def generate(self):
        from hpogrid.idds_interface.steering import SteeringIDDS
        SteeringIDDS()
    

if __name__ == '__main__':
    HPOGridCLI()