"""Train the article's small language model: python train.py [--large]."""
import argparse
from dataclasses import asdict, dataclass
import hashlib
import json
from pathlib import Path

import tensorflow as tf
from datasets import load_dataset

REVISION = "e63b8adc3b1a1bdc7cac5b500d150b71346b0628"


@dataclass(frozen=True)
class Config:
    context: int = 512
    width: int = 512
    layers: int = 6
    heads: int = 8
    batch: int = 32
    steps: int = 20000
    learning_rate: float = 0.0003
    weight_decay: float = 0.01
    seed: int = 42
    report_every: int = 500
    eval_batches: int = 8
    sample_bytes: int = 160
    temperature: float = 0.8
    top_k: int = 40
    prompts: tuple = ("The ", "Once upon a time", "I think")

    def __post_init__(self):
        for key in ("context", "width", "layers", "heads", "batch", "steps",
                    "report_every", "eval_batches", "sample_bytes"):
            value = getattr(self, key)
            if type(value) is not int or value < 1:
                raise ValueError(f"{key} must be a positive integer")
        if self.width % self.heads:
            raise ValueError("width must be divisible by heads")
        if not 0 <= self.seed < 2**31:
            raise ValueError("seed must be in [0, 2**31)")
        if not 0 < self.learning_rate <= 1 or not 0 <= self.weight_decay <= 1:
            raise ValueError("invalid learning rate or weight decay")
        if not 0 < self.temperature <= 10 or not 1 <= self.top_k <= 256:
            raise ValueError("invalid sampling settings")
        if not self.prompts or any(not isinstance(p, str) or not p for p in self.prompts):
            raise ValueError("prompts must contain nonempty strings")

    def parameters(self):
        # Separate input/output embeddings; learned positions; attention biases;
        # FFN expansion 4x; two pre-norms per block and one final norm.
        d = self.width
        return 512*d + self.context*d + 256 + 2*d + self.layers*(12*d*d + 13*d)

    def as_dict(self):
        return asdict(self)


class Block(tf.keras.layers.Layer):
    def __init__(self, config, **kwargs):
        super().__init__(**kwargs)
        self.norm1 = tf.keras.layers.LayerNormalization(epsilon=1e-5)
        self.attention = tf.keras.layers.MultiHeadAttention(
            num_heads=config.heads, key_dim=config.width // config.heads,
            dropout=0.0, flash_attention=False)
        self.norm2 = tf.keras.layers.LayerNormalization(epsilon=1e-5)
        self.expand = tf.keras.layers.Dense(4 * config.width, activation=tf.nn.gelu)
        self.contract = tf.keras.layers.Dense(config.width)

    def call(self, x, training=False):
        normalized = self.norm1(x)
        x = x + self.attention(normalized, normalized,
                               use_causal_mask=True, training=training)
        return x + self.contract(self.expand(self.norm2(x)))


class LanguageModel(tf.keras.Model):
    def __init__(self, config):
        super().__init__()
        self.context = config.context
        self.tokens = tf.keras.layers.Embedding(256, config.width)
        self.positions = tf.keras.layers.Embedding(config.context, config.width)
        # Explicit named attributes also make every block visible to TF checkpoints.
        self.block_names = tuple(f"block_{i}" for i in range(config.layers))
        for name in self.block_names:
            setattr(self, name, Block(config, name=name))
        self.norm = tf.keras.layers.LayerNormalization(epsilon=1e-5)
        self.output_head = tf.keras.layers.Dense(256, dtype="float32")

    def call(self, ids, training=False):
        length = tf.shape(ids)[1]
        tf.debugging.assert_less_equal(length, self.context)
        x = self.tokens(ids) + self.positions(tf.range(length))[None, :, :]
        for name in self.block_names:
            x = getattr(self, name)(x, training=training)
        return self.output_head(self.norm(x))


    @tf.function(input_signature=[tf.TensorSpec((1, None), tf.int32)])
    def next_logits(self, ids):
        """Reuse one traced inference graph while the prompt length changes."""
        return self(ids, training=False)[:, -1, :]


def make_model(config):
    tf.keras.utils.set_random_seed(config.seed)
    model = LanguageModel(config)
    model(tf.zeros((1, config.context), dtype=tf.int32))
    if model.count_params() != config.parameters():
        raise ValueError("Model differs from parameter accounting")
    return model


def make_optimizer(config, model):
    optimizer = tf.keras.optimizers.AdamW(
        learning_rate=config.learning_rate, weight_decay=config.weight_decay,
        global_clipnorm=1.0)
    optimizer.exclude_from_weight_decay(var_names=["bias", "beta", "gamma"])
    optimizer.build(model.trainable_variables)
    return optimizer


def dataset(stories, context, batch_size, training=False):
    texts = tf.data.Dataset.from_generator(
        lambda: (row["story"] for row in stories),
        output_signature=tf.TensorSpec((), tf.string),
    ).cache()  # RAM, not disk.
    data = texts.map(
        lambda text: tf.io.decode_raw(text + "\n\n", tf.uint8),
        num_parallel_calls=tf.data.AUTOTUNE,
    )
    data = data.unbatch().batch(context + 1, drop_remainder=True)
    if training:
        data = data.shuffle(4096, seed=42)
    data = data.batch(batch_size, drop_remainder=True)
    data = data.map(lambda chunk: (chunk[:, :-1], chunk[:, 1:]))
    return data.prefetch(tf.data.AUTOTUNE)


def sample(model, prompt, length=200):
    tokens = list(prompt.encode())
    for i in range(length):
        context = tf.constant([tokens[-model.context:]], dtype=tf.int32)
        logits = model.next_logits(context) / 0.8
        token = tf.random.stateless_categorical(logits, 1, seed=[42, i])
        tokens.append(int(token[0, 0]))
    return bytes(tokens).decode(errors="replace")


class Progress(tf.keras.callbacks.Callback):
    def __init__(self, checkpoints, target_step=None):
        super().__init__()
        self.checkpoints = checkpoints
        self.target_step = target_step

    def on_train_batch_end(self, batch, logs=None):
        if self.target_step is not None and int(self.model.optimizer.iterations) >= self.target_step:
            self.model.stop_training = True

    def on_epoch_end(self, epoch, logs=None):
        step = int(self.model.optimizer.iterations)
        self.checkpoints.save(checkpoint_number=step)
        print(sample(self.model, "Once upon a time"), flush=True)


def configure_hardware():
    gpus = tf.config.list_physical_devices("GPU")
    if gpus:
        tf.config.set_visible_devices(gpus[0], "GPU")
        tf.config.experimental.set_memory_growth(gpus[0], True)


def story_hash(row):
    text = " ".join(row["story"].split())
    return hashlib.sha256(text.encode()).digest()


def load_data(config):
    stories = load_dataset(
        "SimpleStories/SimpleStories", split="train", revision=REVISION,
        streaming=True, columns=["story"],
    )
    # Keep the validation split used by the existing checkpoint.
    protected = {story_hash(row) for row in stories.take(1100)}

    def held_out(row):
        key = story_hash(row)
        return key not in protected and int.from_bytes(key[:4], "big") % 100 == 0

    validation = dataset(stories.filter(held_out).take(1000), config.context, config.batch)
    training = dataset(stories.filter(lambda row: not held_out(row)),
                       config.context, config.batch, training=True)
    return training, validation


def build_model(config):
    model = make_model(config)
    model.compile(
        optimizer=make_optimizer(config, model),
        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    )
    return model


def train(model, training, validation, config, run, initial_checkpoint=None,
          max_checkpoints=3, extra_callbacks=()):
    run.mkdir(parents=True, exist_ok=True)
    checkpoint = tf.train.Checkpoint(model=model, optimizer=model.optimizer)
    checkpoints = tf.train.CheckpointManager(
        checkpoint, str(run / "checkpoints"), max_to_keep=max_checkpoints,
    )
    saved = checkpoints.latest_checkpoint or initial_checkpoint
    if saved:
        checkpoint.restore(saved).expect_partial().assert_existing_objects_matched()
    if int(model.optimizer.iterations) >= config.steps:
        return
    (run / "config.json").write_text(json.dumps(
        {"model": config.as_dict(), "revision": REVISION,
         "initial_checkpoint": initial_checkpoint}, indent=2,
    ))

    try:
        return model.fit(
            training.repeat(), validation_data=validation,
            steps_per_epoch=config.report_every,
            epochs=(config.steps + config.report_every - 1) // config.report_every,
            initial_epoch=int(model.optimizer.iterations) // config.report_every,
            shuffle=False, verbose=2,
            callbacks=[
                Progress(checkpoints, config.steps),
                tf.keras.callbacks.CSVLogger(str(run / "loss.csv"), append=True),
                *extra_callbacks,
            ],
        )
    finally:
        checkpoints.save(checkpoint_number=int(model.optimizer.iterations))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--large", action="store_true", help="Use the 152M model")
    parser.add_argument("--steps", type=int, default=20000)
    parser.add_argument("--run", type=Path)
    args = parser.parse_args()

    configure_hardware()
    config = (Config(width=1024, layers=12, heads=16,
                     learning_rate=0.0002, steps=args.steps)
              if args.large else Config(steps=args.steps))
    run = args.run or Path("runs/large" if args.large else "runs/small")
    training, validation = load_data(config)
    model = build_model(config)
    train(model, training, validation, config, run)


if __name__ == "__main__":
    main()