"""
inference.py
-------------
End-to-end example: load the model from this repo and turn a landmark
sequence into predicted text.

Requirements:
    - transformer_weights.h5   (this repo)
    - modeling.py, config.json (this repo)
    - character_to_prediction_index.json  (copy this in from your own
      training data / Kaggle dataset — it's the char<->id vocabulary and
      is NOT regenerated automatically here)

Run:
    python inference.py --landmarks path/to/landmarks.npy
"""

import argparse
import json

import numpy as np
import tensorflow as tf

from modeling import build_model, pre_process, FEATURE_COLUMNS


def load_vocab(vocab_path="character_to_prediction_index.json"):
    with open(vocab_path, "r") as f:
        char_to_num = json.load(f)
    char_to_num["P"] = 59  # pad
    char_to_num["<"] = 60  # start
    char_to_num[">"] = 61  # end
    num_to_char = {v: k for k, v in char_to_num.items()}
    return char_to_num, num_to_char


def predict(landmarks_np, model, num_to_char, start_token_idx=60, end_token_idx=61):
    """landmarks_np: np.ndarray of shape (num_frames, len(FEATURE_COLUMNS))"""
    x = tf.constant(landmarks_np, dtype=tf.float32)
    x = pre_process(x)[tf.newaxis, ...]  # add batch dim
    token_ids = model.generate(x, start_token_idx)[0].numpy()

    chars = []
    for idx in token_ids[1:]:  # skip the start token
        if idx == end_token_idx:
            break
        chars.append(num_to_char.get(int(idx), ""))
    return "".join(chars)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--landmarks", required=True, help="Path to a .npy file of shape (frames, %d)" % len(FEATURE_COLUMNS))
    parser.add_argument("--weights", default="transformer_weights.h5")
    parser.add_argument("--vocab", default="character_to_prediction_index.json")
    args = parser.parse_args()

    model = build_model()
    model.load_weights(args.weights)

    _, num_to_char = load_vocab(args.vocab)

    landmarks = np.load(args.landmarks)
    text = predict(landmarks, model, num_to_char)
    print("Predicted phrase:", text)