# Source code for archai.nas.exp_runner

    # Copyright (c) Microsoft Corporation.
    # Licensed under the MIT license.

    from typing import Optional, Type, Tuple
    from abc import ABC, abstractmethod
    import shutil
    import os

    from overrides import EnforceOverrides

    from archai.nas.model_desc_builder import ModelDescBuilder
    from archai.nas.arch_trainer import TArchTrainer
    from archai.common import common
    from archai.common import utils
    from archai.common.config import Config
    from archai.nas.evaluater import Evaluater, EvalResult
    from archai.nas.searcher import Searcher, SearchResult
    from archai.nas.finalizers import Finalizers
    from archai.nas.random_finalizers import RandomFinalizers
    from archai.nas.model_desc_builder import ModelDescBuilder

    class ExperimentRunner(ABC, EnforceOverrides):
        def __init__(self, config_filename:str, base_name:str, clean_expdir=False) -> None:
            self.config_filename = config_filename
            self.base_name = base_name
            self.clean_expdir = clean_expdir

        def run_search(self, conf_search:Config)->SearchResult:
            model_desc_builder = self.model_desc_builder()
            trainer_class = self.trainer_class()
            finalizers = self.finalizers()

            search = self.searcher()
            return search.search(conf_search, model_desc_builder, trainer_class, finalizers)

        def run_eval(self, conf_eval:Config)->EvalResult:
            evaler = self.evaluater()
            return evaler.evaluate(conf_eval,
                                   model_desc_builder=self.model_desc_builder())

        def run(self, search=True, eval=True) \
                ->Tuple[Optional[SearchResult], Optional[EvalResult]]:

            search_result, eval_result = None, None

            if search: # run search
                conf = self._init_conf(True, clean_expdir=self.clean_expdir)
                search_result = self.run_search(conf['nas']['search'])

            if eval:
                conf = self.get_conf(False)
                common.clean_ensure_expdir(conf, clean_dir=self.clean_expdir, ensure_dir=True)
                if search:
                    # first copy search result to eval, otherwise we expect eval config to point to results
                    self.copy_search_to_eval()

                conf = self._init_conf(False, clean_expdir=False)
                eval_result = self.run_eval(conf['nas']['eval'])

            return search_result, eval_result

        def copy_search_to_eval(self)->None:
            # do not cache conf_search or conf_eval as it may have values that
            # needs env var expansion.

            # get desc file path that search has produced
            conf_search = self.get_conf(True)['nas']['search']
            search_desc_filename = conf_search['final_desc_filename']
            search_desc_filepath = utils.full_path(search_desc_filename)
            assert search_desc_filepath and os.path.exists(search_desc_filepath)

            # get file path that eval would need
            conf_eval = self.get_conf(False)['nas']['eval']
            eval_desc_filename = conf_eval['final_desc_filename']
            eval_desc_filepath = utils.full_path(eval_desc_filename)
            assert eval_desc_filepath
            utils.copy_file(search_desc_filepath, eval_desc_filepath)

        def model_desc_builder(self)->Optional[ModelDescBuilder]:
            return ModelDescBuilder() # default model desc builder puts nodes with no edges

        def searcher(self)->Searcher:
            return Searcher()

        def evaluater(self)->Evaluater:
            return Evaluater()

        @abstractmethod
        def trainer_class(self)->TArchTrainer:
            pass

        def finalizers(self)->Finalizers:
            conf = common.get_conf()
            finalizer = conf['nas']['search']['finalizer']

            if not finalizer or finalizer == 'default':
                return Finalizers()
            elif finalizer == 'random':
                return RandomFinalizers()
            else:
                raise NotImplementedError

        def get_expname(self, is_search_or_eval:bool)->str:
            return self.base_name + ('_search' if is_search_or_eval else '_eval')

        def get_conf(self, is_search_or_eval:bool)->Config:
            conf = common.create_conf(config_filepath=self.config_filename,
                param_args=['--common.experiment_name', self.get_expname(is_search_or_eval)])
            common.update_envvars(conf) # config paths might include env vars
            return conf

        def _init_conf(self, is_search_or_eval:bool, clean_expdir:bool)->Config:
            config_filename = self.config_filename

            conf = common.common_init(config_filepath=config_filename,
                param_args=['--common.experiment_name', self.get_expname(is_search_or_eval),
                            ], clean_expdir=clean_expdir)
            return conf