import ray.rllib.agents.a3c as a3c
import ray.rllib.agents.ddpg as ddpg
import ray.rllib.agents.dqn as dqn
import ray.rllib.agents.impala as impala
import ray.rllib.agents.ppo as ppo


# import ..agents.ppo as ppo
# import ..agents.a3c as a3c
# import ..agents.impala as impala


def init_trainer(name, config):
    trainer = ''
    if name == 'a3c':
        trainer = a3c.A3CTrainer(config=config)
    elif name == 'impala':
        trainer = impala.ImpalaTrainer(config=config)
    elif name == 'ppo':
        trainer = ppo.PPOTrainer(config=config)
    elif name == 'dqn':
        trainer = dqn.DQNTrainer(config=config)
    elif name == 'ddpg':
        trainer = ddpg.DDPGTrainer(config=config)
    else:
        print('wrong trainer name:{}'.format(name))
        raise NotImplementedError('wrong trainer name:{}'.format(name))
    return trainer


def set_weights(trainer, weights):
    trainer.set_weights(weights)


def get_weights(trainer, policy="ppo_policy"):
    weights = trainer.get_weights([policy])
    return weights


def get_default_config(name):
    if name == 'a3c':
        config = a3c.DEFAULT_CONFIG.copy()
    elif name == 'impala':
        config = impala.DEFAULT_CONFIG.copy()
    elif name == 'ppo':
        config = ppo.DEFAULT_CONFIG.copy()
    elif name == 'dqn':
        config = dqn.DEFAULT_CONFIG.copy()
    elif name == 'ddpg':
        config = ddpg.DEFAULT_CONFIG.copy()
    else:
        print('wrong trainer name:{}'.format(name))
        raise NotImplementedError('wrong trainer name:{}'.format(name))
    return config
