from utils import logger, TaskMode, load_dnn_input_record, load_lr_input_record

feeding_index = {'dnn_input': 0, 'lr_input': 1, 'click': 2}


class Dataset(object):
    def train(self, path):
        '''
        Load trainset.
        '''
        logger.info("load trainset from %s" % path)
        mode = TaskMode.create_train()
        return self._parse_creator(path, mode)

    def test(self, path):
        '''
        Load testset.
        '''
        logger.info("load testset from %s" % path)
        mode = TaskMode.create_test()
        return self._parse_creator(path, mode)

    def infer(self, path):
        '''
        Load infer set.
        '''
        logger.info("load inferset from %s" % path)
        mode = TaskMode.create_infer()
        return self._parse_creator(path, mode)

    def _parse_creator(self, path, mode):
        '''
        Parse dataset.
        '''

        def _parse():
            with open(path) as f:
                for line_id, line in enumerate(f):
                    fs = line.strip().split('\t')
                    dnn_input = load_dnn_input_record(fs[0])
                    lr_input = load_lr_input_record(fs[1])
                    if not mode.is_infer():
                        click = [int(fs[2])]
                        yield dnn_input, lr_input, click
                    else:
                        yield dnn_input, lr_input

        return _parse


def load_data_meta(path):
    '''
    load data meta info from path, return (dnn_input_dim, lr_input_dim)
    '''
    with open(path) as f:
        lines = f.read().split('\n')
        err_info = "wrong meta format"
        assert len(lines) == 2, err_info
        assert 'dnn_input_dim:' in lines[0] and 'lr_input_dim:' in lines[
            1], err_info
        res = map(int, [_.split(':')[1] for _ in lines])
        logger.info('dnn input dim: %d' % res[0])
        logger.info('lr input dim: %d' % res[1])
        return res