#!/usr/bin/env python3
"""Download and deterministically tokenize the frozen Tiny Shakespeare data."""
import argparse
import hashlib
import os
import pickle
import urllib.request

import numpy as np


URL = (
    "https://raw.githubusercontent.com/karpathy/char-rnn/master/"
    "data/tinyshakespeare/input.txt"
)
EXPECTED = {
    "input.txt":
        "86c4e6aa9db7c042ec79f339dcb96d42b0075e16b8fc2e86bf0ca57e2dc565ed",
    "train.bin":
        "6ec305602a99ac2802745a134e1f5e33e2231b4855525b00b9aebb730ac2626f",
    "val.bin":
        "d37d30cc0c8327c270d493299c3dca54135f6d5f1c9ef60cda78076e311204b1",
}


def sha256(path):
    digest = hashlib.sha256()
    with open(path, "rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def verify(path, expected):
    actual = sha256(path)
    if actual != expected:
        raise RuntimeError(
            f"hash mismatch for {path}: expected {expected}, found {actual}"
        )


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--out-dir", required=True)
    args = parser.parse_args()
    root = os.path.abspath(args.out_dir)
    os.makedirs(root, exist_ok=True)
    input_path = os.path.join(root, "input.txt")
    if not os.path.isfile(input_path):
        urllib.request.urlretrieve(URL, input_path)
    verify(input_path, EXPECTED["input.txt"])

    with open(input_path, encoding="utf-8") as handle:
        data = handle.read()
    chars = sorted(set(data))
    stoi = {char: index for index, char in enumerate(chars)}
    itos = {index: char for char, index in stoi.items()}
    encoded = np.asarray([stoi[char] for char in data], dtype=np.uint16)
    boundary = int(len(encoded) * 0.9)
    encoded[:boundary].tofile(os.path.join(root, "train.bin"))
    encoded[boundary:].tofile(os.path.join(root, "val.bin"))
    with open(os.path.join(root, "meta.pkl"), "wb") as handle:
        pickle.dump(
            {"vocab_size": len(chars), "itos": itos, "stoi": stoi},
            handle,
        )

    for name, expected in EXPECTED.items():
        verify(os.path.join(root, name), expected)
    print(
        f"prepared Tiny Shakespeare at {root}: "
        f"{boundary} train, {len(encoded) - boundary} validation tokens"
    )


if __name__ == "__main__":
    main()