import sys
from absl import flags
from functools import partial

import tensorflow.compat.v1 as tf

from reward.comparative.data import \
    SPLITS, LOCAL_TFRECORDS_PATH, DATASET_IS_PACKED, get_dataset

def _define_flags():
    flags.DEFINE_string(
        name="dataset_id",
        default=None,
        help="Id of the dataset generated by jsonl_to_tsvs.py"
    )
    return flags.FLAGS

def _serialize_tokens(*values, colnames):
    example_data = {
        colnames[i]: tf.train.Feature(
            int64_list=tf.train.Int64List(value=tokens)
        )
        for i, tokens in enumerate(values)
    }
    example = tf.train.Example(features=tf.train.Features(feature=example_data))
    return example.SerializeToString()

def _tf_serialize_tokens(tf_tokens_dict):
    serialize_with_colnames = partial(
        _serialize_tokens,
        colnames=list(tf_tokens_dict.keys())
    )
    tf_string = tf.py_function(
        func=serialize_with_colnames,
        inp=list(tf_tokens_dict.values()),
        Tout=tf.string
    )
    return tf.reshape(tf_string, ())

if __name__ == "__main__":
    tf.enable_eager_execution()
    assert (not DATASET_IS_PACKED), "Packed dataset not supported"
    FLAGS = _define_flags()
    FLAGS(sys.argv)
    for split in SPLITS:
        shuffled_dataset = get_dataset(
            dataset_id=FLAGS.dataset_id,
            split=split,
            from_local=True,
            from_tfrecords=False,
            stack_answer_pairs=False,
            shuffle_buffer_size=10000
        )
        serialized_dataset = shuffled_dataset.map(_tf_serialize_tokens)
        # Write tfrecords dataset to disk
        tfrecords_path = LOCAL_TFRECORDS_PATH.format(
            dataset_id=FLAGS.dataset_id,
            split=split
        )
        print(f"Writting '{split}' tfrecords")
        tfrecords_writer = tf.data.experimental.TFRecordWriter(tfrecords_path)
        tfrecords_writer.write(serialized_dataset)