# -*- coding: utf-8 -*-
"""Seq2Seq Addition using LSTM Tensorflow .ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1nOdJF4CsNNYJXKrAvy3nW4uoawFBMdmY
"""

from tensorflow import keras
from tensorflow.keras import layers
import numpy as np

"""##### Generating dataset"""

TRAINING_SIZE = 50000
DIGITS = 3
REVERSE = False

MAX_QUESTION_LEN = DIGITS + 1 + DIGITS
MAX_ANSWER_LEN = DIGITS + 1

class CharacterTable:
    def __init__(self, chars):
        self.chars = sorted(set(chars))
        self.c2i = dict((c, idx) for idx, c in enumerate(self.chars))
        self.i2c = dict((idx, c) for idx, c in enumerate(self.chars))

    def encode(self, math_str, seq_len):
        x = np.zeros((seq_len, len(self.chars)))
        for i, c in enumerate(math_str):
            x[i, self.c2i[c]] = 1

        return x
    
    def decode(self, x, calc_argmax=True):
        if calc_argmax:
            x = x.argmax(axis=-1)

        return "".join(self.i2c[i] for i in x)

chars = "0123456789+ "
ctable = CharacterTable(chars)

questions = []
expected = []
seen = set()
print("Generating data...")
while len(questions) < TRAINING_SIZE:
    f = lambda: int(
        "".join(
            np.random.choice(list("0123456789"))
            for i in range(np.random.randint(1, DIGITS + 1))
        )
    )
    a, b = f(), f()
    # Skip any addition questions we've already seen
    # Also skip any such that x+Y == Y+x (hence the sorting).
    key = tuple(sorted((a, b)))
    if key in seen:
        continue
    seen.add(key)
    # Pad the data with spaces such that it is always MAXLEN.
    q = "{}+{}".format(a, b)
    query = q + " " * (MAX_QUESTION_LEN - len(q))
    ans = str(a + b)
    # Answers can be of maximum size DIGITS + 1.
    ans += " " * (DIGITS + 1 - len(ans))
    if REVERSE:
        # Reverse the query, e.g., '12+345  ' becomes '  543+21'. (Note the
        # space used for padding.)
        query = query[::-1]
    questions.append(query)
    expected.append(ans)
print("Total questions:", len(questions))

questions[:5]

expected[:5]

x = np.zeros((len(questions), MAX_QUESTION_LEN, len(chars)), dtype=np.bool)
y = np.zeros((len(questions), MAX_ANSWER_LEN, len(chars)), dtype=np.bool)

for i, sentence in enumerate(questions):
    x[i] = ctable.encode(math_str=sentence, seq_len=MAX_QUESTION_LEN)

for i, sentence in enumerate(expected):
    y[i] = ctable.encode(math_str=sentence, seq_len=MAX_ANSWER_LEN)

# shuffle
indices = np.arange(len(y))
np.random.shuffle(indices)
x = x[indices]
y = y[indices]


split_at = len(x) - len(x) // 10
(x_train, x_val) = x[:split_at], x[split_at:]
(y_train, y_val) = y[:split_at], y[split_at:]

print(x_train.shape, y_train.shape)
print(x_val.shape, y_val.shape)

"""##### Build model"""

num_layers = 3

model = keras.Sequential()
model.add(layers.LSTM(128, input_shape=(MAX_QUESTION_LEN, len(chars))))
model.add(layers.RepeatVector(MAX_ANSWER_LEN))
for _ in range(num_layers):
    model.add(layers.LSTM(128, return_sequences=True))
model.add(layers.Dense(len(chars), activation="softmax"))
model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])
model.summary()

EPOCHS = 100
BATCH_SIZE = 1048

for epoch in range(1, EPOCHS):
    model.fit(
        x_train,
        y_train,
        batch_size=BATCH_SIZE,
        epochs=1,
        validation_data=(x_val, y_val),
    )
    if epoch % 25 == 0:
        print(f"Epoch: {epoch}")
        for i in range(10):
            ind = np.random.randint(0, len(x_val))
            rowx, rowy = x_val[np.array([ind])], y_val[np.array([ind])]
            preds = np.argmax(model.predict(rowx), axis=-1)
    
            q = ctable.decode(rowx[0])
            correct = ctable.decode(rowy[0])
            guess = ctable.decode(preds[0], calc_argmax=False)
    
            print("Q", q[::-1] if REVERSE else q, end=" ")
            print("T", correct, end=" ")
            if correct == guess:
                print("☑ " + guess)
            else:
                print("☒ " + guess)

"""#### Reference:

1. https://keras.io/examples/nlp/addition_rnn/

"""