import tensorflow as tf
from keras import layers


class PositionEmbedding(layers.Layer):
    def __init__(self, block_size, **kwargs):
        super().__init__(**kwargs)
        self.block_size = block_size

    def build(self, input_shape):
        self.pos_emb = self.add_weight(
            name="pos_emb", shape=(self.block_size, input_shape[-1]), initializer="random_normal"
        )

    def call(self, x):
        T = tf.shape(x)[1]
        return x + self.pos_emb[tf.newaxis, :T, :]

    def get_config(self):
        config = super().get_config()
        config["block_size"] = self.block_size
        return config


def build_gpt(vocab_size, block_size, n_layer=4, n_head=4, n_embd=128, dropout=0.1, name="indigo"):
    tokens = tf.keras.Input(shape=(None,), dtype="int64", name="tokens")
    x = layers.Embedding(vocab_size, n_embd, name="tok_emb")(tokens)
    x = PositionEmbedding(block_size, name="pos_emb")(x)
    x = layers.Dropout(dropout)(x)
    for i in range(n_layer):
        xn = layers.LayerNormalization(epsilon=1e-5, name=f"ln1_{i}")(x)
        attn = layers.MultiHeadAttention(
            num_heads=n_head, key_dim=n_embd // n_head, dropout=dropout, name=f"attn_{i}"
        )
        x = x + attn(xn, xn, use_causal_mask=True)
        xn = layers.LayerNormalization(epsilon=1e-5, name=f"ln2_{i}")(x)
        h = layers.Dense(4 * n_embd, activation="gelu", name=f"fc_{i}")(xn)
        h = layers.Dense(n_embd, name=f"proj_{i}")(h)
        h = layers.Dropout(dropout)(h)
        x = x + h
    x = layers.LayerNormalization(epsilon=1e-5, name="ln_f")(x)
    logits = layers.Dense(vocab_size, use_bias=False, name="head")(x)
    return tf.keras.Model(tokens, logits, name=name)


@tf.function(reduce_retracing=True)
def _langkah(model, idx_cond, temperature, top_k):
    logits = model(idx_cond, training=False)[:, -1, :]
    logits = logits / max(temperature, 1e-8)
    if top_k is not None:
        k = min(top_k, int(logits.shape[-1]))
        vals, _ = tf.math.top_k(logits, k=k)
        logits = tf.where(
            logits < vals[:, -1:],
            tf.fill(tf.shape(logits), tf.float32.min),
            logits,
        )
    return tf.random.categorical(logits, num_samples=1, dtype=tf.int64)


def generate(model, idx, max_new_tokens, block_size, temperature=1.0, top_k=None):
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -block_size:]
        next_id = _langkah(model, idx_cond, temperature, top_k)
        idx = tf.concat([idx, next_id], axis=1)
    return idx