import json
import numpy as np


class DssmHistoryReducer:
    def __init__(self, group_size, emb_size, pad_last):
        self.group_size = group_size
        self.emb_size = emb_size
        self.pad_last = pad_last

    def __call__(self, key, recs):
        history = {'embs': []}

        for i, rec in enumerate(recs, 1):
            assert len(rec['emb']) == self.emb_size
            history['embs'].append(' '.join(map(str, rec['emb'])))

            if i % self.group_size == 0:
                yield {
                    'external_id': key['external_id'],
                    'history': json.dumps(history),
                    'random': np.random.rand(),
                    'fake_category': '0'
                }
                history = {'embs': []}

        if len(history['embs']) > 0:
            if self.pad_last:
                for _ in range(self.group_size - len(history['embs'])):
                    history['embs'].append(' '.join(['0' for __ in range(self.emb_size)]))

            yield {
                'external_id': key['external_id'],
                'history': json.dumps(history),
                'random': np.random.rand(),
                'fake_category': '0'
            }
