# -*- coding: utf-8 -*-
from datacloud.dev_utils.yt import yt_utils
import yt.wrapper as yt_wrapper


@yt_wrapper.with_context
def match_score(_, recs, context):
    score = None
    for rec in recs:
        if context.table_index == 0:
            score = rec['score']
        elif score:
            yield {'hashed_cid': rec['hashed_cid'], 'hashed_id': rec['hashed_id'], 'score': score}


def build_meta_score_table_mapper(rec):
    yield {'hashed_id': rec['hashed_id'], 'score': rec['score']}
    yield {'hashed_id': rec['hashed_cid'], 'score': rec['score']}


def build_meta_score_table(yt_client, small_crypta_table, encoded_score, result_table):
    yt_client.run_sort(
        small_crypta_table,
        sort_by='hashed_cid'
    )
    yt_client.run_sort(
        encoded_score,
        sort_by='hashed_cid'
    )
    tmp_table = result_table + 'tmp'
    yt_client.run_reduce(
        match_score,
        [encoded_score, small_crypta_table],
        yt_client.TablePath(tmp_table, schema=[
            {'name': 'hashed_cid', 'type': 'uint64'},
            {'name': 'hashed_id', 'type': 'uint64'},
            {'name': 'score', 'type': 'double'}]),
        reduce_by='hashed_cid'
    )
    yt_client.run_map(
        build_meta_score_table_mapper,
        tmp_table,
        yt_client.TablePath(
            result_table,
            schema=[{'name': 'hashed_id', 'type': 'uint64'},
                    {'name': 'score', 'type': 'double'}])
    )
    yt_client.remove(tmp_table)


if __name__ == '__main__':
    yt_client = yt_utils.get_yt_client()
    small_crypta_table = '//projects/scoring/tmp/re9ulusv/tables-to-transfer-ydb/mkb/small-crypta-table'

    suffix = ('m1', 'm3', 'm5')
    for suff in suffix:
        print('Start: {}'.format(suff))
        encoded_score = '//projects/scoring/tmp/re9ulusv/tables-to-transfer-ydb/mkb/encoded_score_{}'.format(suff)
        result_table = '//projects/scoring/tmp/re9ulusv/tables-to-transfer-ydb/mkb/encoded_meta_score_{}'.format(suff)
        build_meta_score_table(yt_client, small_crypta_table, encoded_score, result_table)
