import os

import tensorflow as tf

from common import Common

TARGET_INDEX_KEY = "TARGET_INDEX_KEY"
TARGET_STRING_KEY = "TARGET_STRING_KEY"
TARGET_LENGTH_KEY = "TARGET_LENGTH_KEY"
PATH_SOURCE_INDICES_KEY = "PATH_SOURCE_INDICES_KEY"
NODE_INDICES_KEY = "NODES_INDICES_KEY"
PATH_TARGET_INDICES_KEY = "PATH_TARGET_INDICES_KEY"
VALID_CONTEXT_MASK_KEY = "VALID_CONTEXT_MASK_KEY"
PATH_SOURCE_LENGTHS_KEY = "PATH_SOURCE_LENGTHS_KEY"
PATH_LENGTHS_KEY = "PATH_LENGTHS_KEY"
PATH_TARGET_LENGTHS_KEY = "PATH_TARGET_LENGTHS_KEY"
PATH_SOURCE_STRINGS_KEY = "PATH_SOURCE_STRINGS_KEY"
PATH_STRINGS_KEY = "PATH_STRINGS_KEY"
PATH_TARGET_STRINGS_KEY = "PATH_TARGET_STRINGS_KEY"


class Reader:
    class_subtoken_table = None
    class_target_table = None
    class_node_table = None

    def __init__(
        self,
        subtoken_to_index,
        target_to_index,
        node_to_index,
        config,
        is_evaluating=False,
    ):
        tf.logging.set_verbosity(tf.logging.ERROR)
        self.config = config
        self.file_path = (
            config.TEST_PATH if is_evaluating else (config.TRAIN_PATH + ".train.c2s")
        )
        if self.file_path is not None and not os.path.exists(self.file_path):
            print(
                "%s cannot find file: %s"
                % (
                    "Evaluation reader" if is_evaluating else "Train reader",
                    self.file_path,
                )
            )
        self.batch_size = config.TEST_BATCH_SIZE if is_evaluating else config.BATCH_SIZE
        self.is_evaluating = is_evaluating

        self.context_pad = "{},{},{}".format(Common.PAD, Common.PAD, Common.PAD)
        self.record_defaults = [[self.context_pad]] * (
            self.config.DATA_NUM_CONTEXTS + 1
        )

        self.subtoken_table = Reader.get_subtoken_table(subtoken_to_index)
        self.target_table = Reader.get_target_table(target_to_index)
        self.node_table = Reader.get_node_table(node_to_index)
        if self.file_path is not None:
            self.output_tensors = self.compute_output()

    @classmethod
    def get_subtoken_table(cls, subtoken_to_index):
        if cls.class_subtoken_table is None:
            cls.class_subtoken_table = cls.initialize_hash_map(
                subtoken_to_index, subtoken_to_index[Common.UNK]
            )
        return cls.class_subtoken_table

    @classmethod
    def get_target_table(cls, target_to_index):
        if cls.class_target_table is None:
            cls.class_target_table = cls.initialize_hash_map(
                target_to_index, target_to_index[Common.UNK]
            )
        return cls.class_target_table

    @classmethod
    def get_node_table(cls, node_to_index):
        if cls.class_node_table is None:
            cls.class_node_table = cls.initialize_hash_map(
                node_to_index, node_to_index[Common.UNK]
            )
        return cls.class_node_table

    @classmethod
    def initialize_hash_map(cls, word_to_index, default_value):
        return tf.contrib.lookup.HashTable(  # TODO: Refactor tf.contrib
            tf.contrib.lookup.KeyValueTensorInitializer(  # TODO: Refactor tf.contrib
                list(word_to_index.keys()),
                list(word_to_index.values()),
                key_dtype=tf.string,
                value_dtype=tf.int32,
            ),
            default_value,
        )

    def process_from_placeholder(self, row):
        parts = tf.io.decode_csv(
            row,
            record_defaults=self.record_defaults,
            field_delim=" ",
            use_quote_delim=False,
        )
        return self.process_dataset(*parts)

    def process_dataset(self, *row_parts):
        row_parts = list(row_parts)
        word = row_parts[0]  # (, )

        if not self.is_evaluating and self.config.RANDOM_CONTEXTS:
            all_contexts = tf.stack(row_parts[1:])
            all_contexts_padded = tf.concat([all_contexts, [self.context_pad]], axis=-1)
            index_of_blank_context = tf.where(
                tf.equal(all_contexts_padded, self.context_pad)
            )
            num_contexts_per_example = tf.reduce_min(index_of_blank_context)

            # if there are less than self.max_contexts valid contexts, still sample self.max_contexts
            safe_limit = tf.cast(
                tf.maximum(num_contexts_per_example, self.config.MAX_CONTEXTS), tf.int32
            )
            rand_indices = tf.random_shuffle(tf.range(safe_limit))[
                : self.config.MAX_CONTEXTS
            ]
            contexts = tf.gather(all_contexts, rand_indices)  # (max_contexts,)
        else:
            contexts = row_parts[1 : (self.config.MAX_CONTEXTS + 1)]  # (max_contexts,)

        # contexts: (max_contexts, )
        split_contexts = tf.string_split(contexts, delimiter=",", skip_empty=False)
        sparse_split_contexts = tf.sparse.SparseTensor(
            indices=split_contexts.indices,
            values=split_contexts.values,
            dense_shape=[self.config.MAX_CONTEXTS, 3],
        )
        dense_split_contexts = tf.reshape(
            tf.sparse.to_dense(
                sp_input=sparse_split_contexts, default_value=Common.PAD
            ),
            shape=[self.config.MAX_CONTEXTS, 3],
        )  # (batch, max_contexts, 3)

        split_target_labels = tf.string_split(tf.expand_dims(word, -1), delimiter="|")
        target_dense_shape = [
            1,
            tf.maximum(
                tf.cast(self.config.MAX_TARGET_PARTS, tf.int64),
                #                tf.to_int64(self.config.MAX_TARGET_PARTS),
                split_target_labels.dense_shape[1] + 1,
            ),
        ]
        sparse_target_labels = tf.sparse.SparseTensor(
            indices=split_target_labels.indices,
            values=split_target_labels.values,
            dense_shape=target_dense_shape,
        )
        dense_target_label = tf.reshape(
            tf.sparse.to_dense(sp_input=sparse_target_labels, default_value=Common.PAD),
            [-1],
        )
        index_of_blank = tf.where(tf.equal(dense_target_label, Common.PAD))
        target_length = tf.reduce_min(index_of_blank)
        dense_target_label = dense_target_label[: self.config.MAX_TARGET_PARTS]
        clipped_target_lengths = tf.clip_by_value(
            target_length, clip_value_min=0, clip_value_max=self.config.MAX_TARGET_PARTS
        )
        target_word_labels = tf.concat(
            [self.target_table.lookup(dense_target_label), [0]], axis=-1
        )  # (max_target_parts + 1) of int

        path_source_strings = tf.slice(
            dense_split_contexts, [0, 0], [self.config.MAX_CONTEXTS, 1]
        )  # (max_contexts, 1)
        flat_source_strings = tf.reshape(path_source_strings, [-1])  # (max_contexts)
        split_source = tf.string_split(
            flat_source_strings, delimiter="|", skip_empty=False
        )  # (max_contexts, max_name_parts)

        sparse_split_source = tf.sparse.SparseTensor(
            indices=split_source.indices,
            values=split_source.values,
            dense_shape=[
                self.config.MAX_CONTEXTS,
                tf.maximum(
                    tf.to_int64(self.config.MAX_NAME_PARTS), split_source.dense_shape[1]
                ),
            ],
        )
        dense_split_source = tf.sparse.to_dense(
            sp_input=sparse_split_source, default_value=Common.PAD
        )  # (max_contexts, max_name_parts)
        dense_split_source = tf.slice(
            dense_split_source, [0, 0], [-1, self.config.MAX_NAME_PARTS]
        )
        path_source_indices = self.subtoken_table.lookup(
            dense_split_source
        )  # (max_contexts, max_name_parts)
        path_source_lengths = tf.reduce_sum(
            tf.cast(tf.not_equal(dense_split_source, Common.PAD), tf.int32), -1
        )  # (max_contexts)

        path_strings = tf.slice(
            dense_split_contexts, [0, 1], [self.config.MAX_CONTEXTS, 1]
        )
        flat_path_strings = tf.reshape(path_strings, [-1])
        split_path = tf.string_split(flat_path_strings, delimiter="|", skip_empty=False)
        sparse_split_path = tf.sparse.SparseTensor(
            indices=split_path.indices,
            values=split_path.values,
            dense_shape=[self.config.MAX_CONTEXTS, self.config.MAX_PATH_LENGTH],
        )
        dense_split_path = tf.sparse.to_dense(
            sp_input=sparse_split_path, default_value=Common.PAD
        )  # (batch, max_contexts, max_path_length)

        node_indices = self.node_table.lookup(
            dense_split_path
        )  # (max_contexts, max_path_length)
        path_lengths = tf.reduce_sum(
            tf.cast(tf.not_equal(dense_split_path, Common.PAD), tf.int32), -1
        )  # (max_contexts)

        path_target_strings = tf.slice(
            dense_split_contexts, [0, 2], [self.config.MAX_CONTEXTS, 1]
        )  # (max_contexts, 1)
        flat_target_strings = tf.reshape(path_target_strings, [-1])  # (max_contexts)
        split_target = tf.string_split(
            flat_target_strings, delimiter="|", skip_empty=False
        )  # (max_contexts, max_name_parts)
        sparse_split_target = tf.sparse.SparseTensor(
            indices=split_target.indices,
            values=split_target.values,
            dense_shape=[
                self.config.MAX_CONTEXTS,
                tf.maximum(
                    tf.to_int64(self.config.MAX_NAME_PARTS), split_target.dense_shape[1]
                ),
            ],
        )
        dense_split_target = tf.sparse.to_dense(
            sp_input=sparse_split_target, default_value=Common.PAD
        )  # (max_contexts, max_name_parts)
        dense_split_target = tf.slice(
            dense_split_target, [0, 0], [-1, self.config.MAX_NAME_PARTS]
        )
        path_target_indices = self.subtoken_table.lookup(
            dense_split_target
        )  # (max_contexts, max_name_parts)
        path_target_lengths = tf.reduce_sum(
            tf.cast(tf.not_equal(dense_split_target, Common.PAD), tf.int32), -1
        )  # (max_contexts)

        valid_contexts_mask = tf.cast(  # tf.to_float
            tf.not_equal(
                tf.reduce_max(path_source_indices, -1)
                + tf.reduce_max(node_indices, -1)
                + tf.reduce_max(path_target_indices, -1),
                0,
            ),
            dtype=tf.float32,
        )

        return {
            TARGET_STRING_KEY: word,
            TARGET_INDEX_KEY: target_word_labels,
            TARGET_LENGTH_KEY: clipped_target_lengths,
            PATH_SOURCE_INDICES_KEY: path_source_indices,
            NODE_INDICES_KEY: node_indices,
            PATH_TARGET_INDICES_KEY: path_target_indices,
            VALID_CONTEXT_MASK_KEY: valid_contexts_mask,
            PATH_SOURCE_LENGTHS_KEY: path_source_lengths,
            PATH_LENGTHS_KEY: path_lengths,
            PATH_TARGET_LENGTHS_KEY: path_target_lengths,
            PATH_SOURCE_STRINGS_KEY: path_source_strings,
            PATH_STRINGS_KEY: path_strings,
            PATH_TARGET_STRINGS_KEY: path_target_strings,
        }

    def reset(self, sess):
        sess.run(self.reset_op)

    def get_output(self):
        return self.output_tensors

    def compute_output(self):
        dataset = tf.data.experimental.CsvDataset(
            self.file_path,
            record_defaults=self.record_defaults,
            field_delim=" ",
            use_quote_delim=False,
            buffer_size=self.config.CSV_BUFFER_SIZE,
        )

        if not self.is_evaluating:
            if self.config.SAVE_EVERY_EPOCHS > 1:
                dataset = dataset.repeat(self.config.SAVE_EVERY_EPOCHS)
            dataset = dataset.shuffle(
                self.config.SHUFFLE_BUFFER_SIZE, reshuffle_each_iteration=True
            )
        dataset = dataset.apply(
            tf.data.experimental.map_and_batch(
                map_func=self.process_dataset,
                batch_size=self.batch_size,
                num_parallel_batches=self.config.READER_NUM_PARALLEL_BATCHES,
            )
        )
        dataset = dataset.prefetch(
            tf.contrib.data.AUTOTUNE
        )  # TODO: Refactor tf.contrib
        self.iterator = dataset.make_initializable_iterator()
        self.reset_op = self.iterator.initializer
        return self.iterator.get_next()


if __name__ == "__main__":
    target_word_to_index = {
        Common.PAD: 0,
        Common.UNK: 1,
        Common.SOS: 2,
        "a": 3,
        "b": 4,
        "c": 5,
        "d": 6,
        "t": 7,
    }
    subtoken_to_index = {Common.PAD: 0, Common.UNK: 1, "a": 2, "b": 3, "c": 4, "d": 5}
    node_to_index = {Common.PAD: 0, Common.UNK: 1, "1": 2, "2": 3, "3": 4, "4": 5}
    import numpy as np

    class Config:
        def __init__(self):
            self.SAVE_EVERY_EPOCHS = 1
            self.TRAIN_PATH = self.TEST_PATH = "test_input/test_input"
            self.BATCH_SIZE = 2
            self.TEST_BATCH_SIZE = self.BATCH_SIZE
            self.READER_NUM_PARALLEL_BATCHES = 1
            self.READING_BATCH_SIZE = 2
            self.SHUFFLE_BUFFER_SIZE = 100
            self.MAX_CONTEXTS = 4
            self.DATA_NUM_CONTEXTS = 4
            self.MAX_PATH_LENGTH = 3
            self.MAX_NAME_PARTS = 2
            self.MAX_TARGET_PARTS = 4
            self.RANDOM_CONTEXTS = True
            self.CSV_BUFFER_SIZE = None

    config = Config()
    reader = Reader(
        subtoken_to_index, target_word_to_index, node_to_index, config, False
    )

    output = reader.get_output()
    target_index_op = output[TARGET_INDEX_KEY]
    target_string_op = output[TARGET_STRING_KEY]
    target_length_op = output[TARGET_LENGTH_KEY]
    path_source_indices_op = output[PATH_SOURCE_INDICES_KEY]
    node_indices_op = output[NODE_INDICES_KEY]
    path_target_indices_op = output[PATH_TARGET_INDICES_KEY]
    valid_context_mask_op = output[VALID_CONTEXT_MASK_KEY]
    path_source_lengths_op = output[PATH_SOURCE_LENGTHS_KEY]
    path_lengths_op = output[PATH_LENGTHS_KEY]
    path_target_lengths_op = output[PATH_TARGET_LENGTHS_KEY]
    path_source_strings_op = output[PATH_SOURCE_STRINGS_KEY]
    path_strings_op = output[PATH_STRINGS_KEY]
    path_target_strings_op = output[PATH_TARGET_STRINGS_KEY]

    sess = tf.InteractiveSession()
    tf.group(
        tf.global_variables_initializer(),
        tf.local_variables_initializer(),
        tf.tables_initializer(),
    ).run()
    reader.reset(sess)

    try:
        while True:
            target_indices, target_strings, target_lengths, path_source_indices, node_indices, path_target_indices, valid_context_mask, path_source_lengths, path_lengths, path_target_lengths, path_source_strings, path_strings, path_target_strings = sess.run(
                [
                    target_index_op,
                    target_string_op,
                    target_length_op,
                    path_source_indices_op,
                    node_indices_op,
                    path_target_indices_op,
                    valid_context_mask_op,
                    path_source_lengths_op,
                    path_lengths_op,
                    path_target_lengths_op,
                    path_source_strings_op,
                    path_strings_op,
                    path_target_strings_op,
                ]
            )

            print("Target strings: ", Common.binary_to_string_list(target_strings))
            print(
                "Context strings: ",
                Common.binary_to_string_3d(
                    np.concatenate(
                        [path_source_strings, path_strings, path_target_strings], -1
                    )
                ),
            )
            print("Target indices: ", target_indices)
            print("Target lengths: ", target_lengths)
            print(
                "Path source strings: ", Common.binary_to_string_3d(path_source_strings)
            )
            print("Path source indices: ", path_source_indices)
            print("Path source lengths: ", path_source_lengths)
            print("Path strings: ", Common.binary_to_string_3d(path_strings))
            print("Node indices: ", node_indices)
            print("Path lengths: ", path_lengths)
            print(
                "Path target strings: ", Common.binary_to_string_3d(path_target_strings)
            )
            print("Path target indices: ", path_target_indices)
            print("Path target lengths: ", path_target_lengths)
            print("Valid context mask: ", valid_context_mask)

    except tf.errors.OutOfRangeError:
        print("Done training, epoch reached")
