import time
import os
from ding.interaction import Slave, TaskFail
from ding.utils import lists_to_dicts


class NaiveLearner(Slave):

    def __init__(self, *args, prefix='', **kwargs):
        super().__init__(*args, **kwargs)
        self._prefix = prefix

    def _process_task(self, task):
        task_name = task['name']
        if task_name == 'resource':
            return {'cpu': 'xxx', 'gpu': 'xxx'}
        elif task_name == 'learner_start_task':
            time.sleep(1)
            self.task_info = task['task_info']
            self.count = 0
            return {'message': 'learner task has started'}
        elif task_name == 'learner_get_data_task':
            time.sleep(0.01)
            return {
                'task_id': self.task_info['task_id'],
                'buffer_id': self.task_info['buffer_id'],
                'batch_size': 2,
                'cur_learner_iter': 1
            }
        elif task_name == 'learner_learn_task':
            data = task['data']
            if data is None:
                raise TaskFail(result={'message': 'no data'})
            time.sleep(0.1)
            data = lists_to_dicts(data)
            assert 'data_id' in data.keys()
            priority_keys = ['replay_unique_id', 'replay_buffer_idx', 'priority']
            self.count += 1
            ret = {
                'info': {
                    'learner_step': self.count
                },
                'task_id': self.task_info['task_id'],
                'buffer_id': self.task_info['buffer_id']
            }
            ret['info']['priority_info'] = {k: data[k] for k in priority_keys}
            if self.count > 5:
                ret['info']['learner_done'] = True
                os.popen('touch {}_final_model.pth'.format(self._prefix))
            return ret
        elif task_name == 'learner_close_task':
            return {'task_id': self.task_info['task_id'], 'buffer_id': self.task_info['buffer_id']}
        else:
            raise TaskFail(
                result={'message': 'task name error'}, message='illegal collector task <{}>'.format(task_name)
            )