import logging
import hydra
from omegaconf import DictConfig, OmegaConf
import argparse
from src import create_logger, copy_all_src
import os

@hydra.main(config_path=None, config_name=None)
def main(cfg: DictConfig):
    print(OmegaConf.to_yaml(cfg))

    env_params = cfg.env_params
    tester_params = cfg.tester_params
    logger_params = cfg.logger_params

    create_logger(logger_params)
    _print_config(env_params, tester_params, logger_params)

    # Import the Tester based on the mode in tester_params
    if tester_params['mode'] == "SA":
        from src.search_sa import Search
    elif tester_params['mode'] == "SA2":
        from src.search_sa_2 import Search
    elif tester_params['mode'] == "SA_AGENT":
        from src.search_sa_agent import Search
    elif tester_params['mode'] == "ILS":
        raise NotImplementedError

    tester = Search(env_params=env_params,
                    tester_params=tester_params)

    tester.run()

def _print_config(env_params, tester_params, logger_params):
    logger = logging.getLogger('root')
    logger.info(f'env_params: {env_params}')
    logger.info(f'tester_params: {tester_params}')
    logger.info(f'logger_params: {logger_params}')

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Test model with config file")
    parser.add_argument('config', type=str, help="Path to the config file")
    parser.add_argument('overrides', nargs='*', help="Override configuration parameters (key=value)")

    args = parser.parse_args()

    # Always make the config path relative to 'configs/eval'
    full_config_path = os.path.join("configs/eval", args.config)

    # Split the path into directory and filename
    config_path, config_name = os.path.split(full_config_path)

    # Initialize Hydra with the provided configuration file
    hydra.initialize(config_path=config_path, version_base=None)
    cfg = hydra.compose(config_name=config_name, overrides=args.overrides)

    main(cfg)
