from ray.rllib.agents.ppo import PPOTFPolicy, PPOTorchPolicy
from ray.rllib.agents.a3c.a3c_tf_policy import A3CTFPolicy
from ray.rllib.agents.a3c.a3c_torch_policy import A3CTorchPolicy
from ray.rllib.agents.dqn import DQNTFPolicy, DQNTorchPolicy
from ray.rllib.agents.ddpg.ddpg_tf_policy import DDPGTFPolicy
from ray.rllib.agents.ddpg.ddpg_torch_policy import DDPGTorchPolicy
from ray.rllib.agents.impala.vtrace_tf_policy import VTraceTFPolicy
from ray.rllib.agents.impala.vtrace_torch_policy import VTraceTorchPolicy


def init_policy(name='ppo', frame='torch'):
    policy = ''
    # if name == 'a3c':
    #     if frame == 'torch':
    #         policy = A3CTorchPolicy
    #     elif frame == 'tf':
    #         policy = A3CTFPolicy
    if name == 'ppo':
        if frame == 'torch':
            policy = PPOTorchPolicy
        elif frame == 'tf':
            policy = PPOTFPolicy
    elif name == 'dqn':
        if frame == 'torch':
            policy = DQNTorchPolicy
        elif frame == 'tf':
            policy = DQNTFPolicy
    elif name == 'ddpg':
        if frame == 'torch':
            policy = DDPGTorchPolicy
        elif frame == 'tf':
            policy = DDPGTFPolicy
    elif name == 'impala':
        if frame == 'torch':
            policy = VTraceTorchPolicy
        elif frame == 'tf':
            policy = VTraceTFPolicy
    else:
        print('wrong policy name:{}'.format(name))
        raise NotImplementedError('wrong policy name:{}'.format(name))
    return policy
