import json
from argparse import ArgumentParser
from typing import List, Tuple

import tensorflow as tf


class SpacingModel(tf.keras.Model):
    def __init__(
            self,
            vocab_size: int,
            hidden_size: int,
            num_classes: int = 3,
            conv_activation: str = "relu",
            dense_activation: str = "relu",
            conv_kernel_and_filter_sizes: List[Tuple[int, int]] = [
                (2, 8),
                (3, 8),
                (4, 8),
                (5, 8),
            ],
            dropout_rate: float = 0.3,
    ):
        super().__init__()

        self.embeddings = tf.keras.layers.Embedding(vocab_size, hidden_size)
        self.convs = [
            tf.keras.layers.Conv1D(
                filter_size,
                kernel_size,
                padding="same",
                activation=conv_activation,
            )
            for kernel_size, filter_size in conv_kernel_and_filter_sizes
        ]
        self.pools = [
            tf.keras.layers.MaxPooling1D(pool_size=filter_size, data_format="channels_first")
            for _, filter_size in conv_kernel_and_filter_sizes
        ]
        self.dropout1 = tf.keras.layers.Dropout(rate=dropout_rate)
        self.output_dense1 = tf.keras.layers.Dense(hidden_size, activation=dense_activation)
        self.dropout2 = tf.keras.layers.Dropout(rate=dropout_rate)
        self.output_dense2 = tf.keras.layers.Dense(num_classes)

    def call(self, input_tensor):
        """
        input_tensor: Tokenized Sequences, Shape: (Batch Size, Sequence Length)
        """

        # embeddings: (Batch Size, Sequence Length, Hidden Size)
        embeddings = self.embeddings(input_tensor)
        # features: (Batch Size, Sequence Length, sum(#filters))
        features = self.dropout1(
            tf.concat([pool(conv(embeddings)) for conv, pool in zip(self.convs, self.pools)], axis=-1)
        )
        # projected: (Batch Size, Sequence Length, Hidden Size)
        projected = self.dropout2(self.output_dense1(features))
        # (Batch Size, Sequence Length, 2)
        return self.output_dense2(projected)


def string_to_example(
        vocab_table: tf.lookup.StaticHashTable,
        encoding: str = "UTF-8",
        max_length: int = 256,
        delete_prob: float = 0.5,
        add_prob: float = 0.15,
):
    @tf.function
    def _inner(tensors: tf.Tensor):
        bytes_array = tf.strings.unicode_split(tf.strings.regex_replace(tensors, " +", " "), encoding)
        space_positions = bytes_array == " "
        sequence_length = tf.shape(space_positions)[0]

        while_condition = lambda i, *_: i < sequence_length

        def while_body(i, strings, labels):
            # 다음 char가 space가 아니고, 문장 끝이 아닐 때 add_prob의 확률로 space 추가
            # 이번 char가 space일 때
            is_next_char_space = tf.cond(i < sequence_length - 1, lambda: bytes_array[i + 1] == " ", lambda: False)

            state = tf.cond(
                is_next_char_space,
                lambda: tf.cond(tf.random.uniform([]) < delete_prob, lambda: 2, lambda: 0),
                lambda: tf.cond(bytes_array[i] != " " and tf.random.uniform([]) < add_prob, lambda: 1, lambda: 0),
            )
            # 0: 그대로 진행
            # 1: 다음 인덱스에 space 추가
            # 2: 다음 space 삭제
            strings = tf.cond(
                state != 1,
                lambda: tf.concat([strings, [bytes_array[i]]], axis=0),
                lambda: tf.concat([strings, [bytes_array[i], " "]], axis=0),
            )
            # label 0: 변화 x
            # label 1: 다음 인덱스에 space 추가
            # label 2: 현재 space 삭제
            labels = tf.cond(
                state == 0,
                lambda: tf.concat([labels, [0]], axis=0),
                lambda: tf.cond(
                    state == 1,
                    lambda: tf.concat([labels, [0, 2]], axis=0),
                    lambda: tf.concat([labels, [1]], axis=0),
                ),
            )
            i += tf.cond(state == 2, lambda: 2, lambda: 1)

            return (i, strings, labels)

        i, strings, labels = tf.while_loop(
            while_condition,
            while_body,
            (
                tf.constant(0),
                tf.constant([], dtype=tf.string),
                tf.constant([], dtype=tf.int32),
            ),
            shape_invariants=(tf.TensorShape([]), tf.TensorShape([None]), tf.TensorShape([None])),
        )

        strings = vocab_table.lookup(tf.concat([["<s>"], strings, ["</s>"]], axis=0))
        labels = tf.concat([[0], labels, [0]], axis=0)

        strings = tf.cond(tf.shape(strings)[0] > max_length, lambda: strings[:max_length], lambda: strings)
        labels = tf.cond(tf.shape(labels)[0] > max_length, lambda: labels[:max_length], lambda: labels)

        length_to_pad = max_length - tf.shape(strings)[0]
        strings = tf.pad(strings, [[0, length_to_pad]])
        labels = tf.pad(labels, [[0, length_to_pad]], constant_values=-1)

        return (strings, labels)

    return _inner


def sparse_categorical_crossentropy_with_ignore(y_true, y_pred, from_logits=False, axis=-1, ignore_id=-1):
    positions = tf.where(y_true != ignore_id)

    y_true = tf.gather_nd(y_true, positions)
    y_pred = tf.gather_nd(y_pred, positions)

    return tf.keras.losses.sparse_categorical_crossentropy(y_true, y_pred, from_logits=from_logits, axis=axis)


def sparse_categorical_accuracy_with_ignore(y_true, y_pred, ignore_id=-1):
    positions = tf.where(y_true != ignore_id)

    y_true = tf.gather_nd(y_true, positions)
    y_pred = tf.gather_nd(y_pred, positions)

    return tf.keras.metrics.sparse_categorical_accuracy(y_true, y_pred)


class SparseCategoricalCrossentropyWithIgnore(tf.python.keras.losses.LossFunctionWrapper):
    def __init__(
            self,
            from_logits=False,
            reduction=tf.keras.losses.Reduction.AUTO,
            ignore_id=-1,
            name="sparse_categorical_crossentropy_with_ignore",
    ):
        super(SparseCategoricalCrossentropyWithIgnore, self).__init__(
            sparse_categorical_crossentropy_with_ignore,
            name=name,
            reduction=reduction,
            ignore_id=ignore_id,
            from_logits=from_logits,
        )


def train(model_path='./models/checkpoint-{epoch}.ckpt',
          log_dir='./logs',
          train_file='',
          dev_file='',
          training_config='',
          char_file=''):
    # args = parser.parse_args()
    with open(training_config) as f:
        config = json.load(f)

    with open(char_file) as f:
        content = f.read()
        keys = ["<pad>", "<s>", "</s>", "<unk>"] + list(content)
        values = list(range(len(keys)))

        vocab_initializer = tf.lookup.KeyValueTensorInitializer(keys, values, key_dtype=tf.string, value_dtype=tf.int32)
        vocab_table = tf.lookup.StaticHashTable(vocab_initializer, default_value=3)

    train_dataset = (
        tf.data.TextLineDataset(tf.constant([train_file]))
            .shuffle(25000000)
            .map(
            string_to_example(vocab_table),
            num_parallel_calls=tf.data.experimental.AUTOTUNE,
        )
            .batch(config["train_batch_size"])
    )
    dev_dataset = (
        tf.data.TextLineDataset(tf.constant([dev_file]))
            .shuffle(100000)
            .map(
            string_to_example(vocab_table),
            num_parallel_calls=tf.data.experimental.AUTOTUNE,
        )
            .batch(config["val_batch_size"])
            .take(4)
    )

    model = SpacingModel(
        config["vocab_size"],
        config["hidden_size"],
        conv_activation=config["conv_activation"],
        dense_activation=config["dense_activation"],
        conv_kernel_and_filter_sizes=config["conv_kernel_and_filter_sizes"],
        dropout_rate=config["dropout_rate"],
    )
    model.compile(
        optimizer=tf.optimizers.Adam(learning_rate=config["learning_rate"]),
        loss=SparseCategoricalCrossentropyWithIgnore(from_logits=True, ignore_id=-1),
        metrics=[sparse_categorical_accuracy_with_ignore],
    )
    model.fit(
        train_dataset,
        epochs=config["epochs"],
        validation_data=dev_dataset,
        steps_per_epoch=400,
        callbacks=[
            tf.keras.callbacks.ModelCheckpoint(filepath=model_path),
            tf.keras.callbacks.TensorBoard(log_dir=log_dir),
            tf.keras.callbacks.ReduceLROnPlateau(patience=2, verbose=1),
        ],
    )

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--train-file", type=str, required=True)
    parser.add_argument("--dev-file", type=str, required=True)
    parser.add_argument("--training-config", type=str, required=True)
    parser.add_argument("--char-file", type=str, required=True)
    parser.add_argument("--model-path", type=str, required=True)
    parser.add_argument("--log-dir", type=str, required=True)
    args = parser.parse_args()

    model_path = args.model_path
    log_dir = args.log_dir
    train_file = args.train_file
    dev_file = args.dev_file
    training_config = args.training_config
    char_file = args.char_file

    train(model_path=model_path, log_dir=log_dir,
          train_file=train_file, dev_file=dev_file,
          training_config=training_config, char_file=char_file)
