from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

import ray
from ray.rllib.ddpg2.models import DDPGModel
from ray.rllib.models.catalog import ModelCatalog
from ray.rllib.optimizers import PolicyEvaluator
from ray.rllib.utils.filter import NoFilter
from ray.rllib.utils.process_rollout import process_rollout
from ray.rllib.utils.sampler import SyncSampler


class DDPGEvaluator(PolicyEvaluator):

    def __init__(self, registry, env_creator, config):
        self.env = ModelCatalog.get_preprocessor_as_wrapper(
            registry, env_creator(config["env_config"]))

        # contains model, target_model
        self.model = DDPGModel(registry, self.env, config)

        self.sampler = SyncSampler(
                        self.env, self.model.model, NoFilter(),
                        config["num_local_steps"], horizon=config["horizon"])

    def sample(self):
        """Returns a batch of samples."""

        rollout = self.sampler.get_data()
        rollout.data["weights"] = 
np.ones_like(rollout.data["rewards"])