import tensorflow as tf
import numpy as np

# Hyperparameters
batch_size = 96
block_size = 20
max_iters = 200
eval_interval = 100
learning_rate = 1e-3
device = 'cuda:0' if tf.test.is_gpu_available() else 'cpu'
eval_iters = 20  # Changed to 20 for consistency with eval_interval
n_embd = 16
n_head = 2
n_layer = 2
dropout_rate = 0.0  # Renamed from dropout for clarity

hidden_nodes = 20

tf.random.set_seed(1337)

with open('data/universe.txt', 'r', encoding='utf-8') as f:
    text = f.read()

data = np.array([ord(c) for c in text], dtype=np.int32)
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]

@tf.function
def get_batch(split):
    data = train_data if split == 'train' else val_data
    ix = np.random.randint(0, len(data) - block_size, size=(batch_size,))
    x = np.array([data[i:i + block_size] for i in ix])
    y = np.array([data[i + 1:i + block_size + 1] for i in ix])
    return x, y

class GPTLanguageModel(tf.keras.Model):
    def __init__(self, vocab_size, n_embd, block_size, n_head, n_layer, dropout_rate=0.1):
        super(GPTLanguageModel, self).__init__()
        self.token_embedding_table = tf.keras.layers.Embedding(vocab_size, n_embd)
        self.position_embedding_table = tf.keras.layers.Embedding(block_size, n_embd)
        self.blocks = [self._build_block(n_embd, n_head) for _ in range(n_layer)]
        self.ln_f = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.lm_head = tf.keras.layers.Dense(vocab_size)

    def _build_block(self, n_embd, n_head):
        head_size = n_embd // n_head  # Bug fixed, removed '+ 1'
        return tf.keras.Sequential([
            tf.keras.layers.LayerNormalization(epsilon=1e-6),
            tf.keras.layers.MultiHeadAttention(num_heads=n_head, key_dim=head_size),
            tf.keras.layers.Dense(4 * n_embd, activation='relu'),
            tf.keras.layers.Dense(n_embd)
        ])

    def call(self, idx, targets=None, training=False):
        B, T = tf.shape(idx)[0], tf.shape(idx)[1]
        tok_emb = self.token_embedding_table(idx)
        pos_emb = self.position_embedding_table(tf.range(T))
        x = tok_emb + tf.cast(pos_emb[None, :, :], tok_emb.dtype)  # Explicit casting to match dtype

        for block in self.blocks:
            x = block(x, training=training)

        x = self.ln_f(x)
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            loss = self.compute_loss(logits, targets)

        return logits, loss

    def compute_loss(self, logits, targets):
        loss = tf.keras.losses.sparse_categorical_crossentropy(targets, logits, from_logits=True)
        return tf.reduce_mean(loss)

    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            logits, _ = self(idx[:, -block_size:], training=False)
            logits = logits[:, -1, :]
            idx_next = tf.random.categorical(logits, num_samples=1, dtype=tf.int32)
            idx = tf.concat([idx, idx_next], axis=1)
        return idx

vocab_size = len(set(data))

model = GPTLanguageModel(vocab_size, n_embd, block_size, n_head, n_layer, dropout_rate)

optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)

train_loss_metric = tf.keras.metrics.Mean(name='train_loss')
val_loss_metric = tf.keras.metrics.Mean(name='val_loss')

@tf.function
def train_step(inputs, targets):
    with tf.GradientTape() as tape:
        _, loss = model(inputs, targets, training=True)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    train_loss_metric.update_state(loss)

@tf.function
def val_step(inputs, targets):
    _, loss = model(inputs, targets, training=False)
    val_loss_metric.update_state(loss)

# Convert numpy operations to Tensor operations
def numpy_to_tensor(numpy_function):
    def tensor_function(*args, **kwargs):
        tensor_args = [tf.convert_to_tensor(arg) for arg in args]
        return numpy_function(*tensor_args, **kwargs)
    return tensor_function

get_batch = numpy_to_tensor(get_batch)

for iteration in range(max_iters):
    for _ in range(eval_interval):
        inputs, targets = get_batch('train')
        train_step(inputs, targets)

    for _ in range(eval_iters):
        inputs, targets = get_batch('val')
        val_step(inputs, targets)

    if (iteration + 1) % eval_interval == 0:
        train_loss = train_loss_metric.result()
        val_loss = val_loss_metric.result()
        print(f'Iteration {iteration + 1}: Train Loss = {train_loss:.4f}, Val Loss = {val_loss:.4f}')
        train_loss_metric.reset_states()
        val_loss_metric.reset_states()

model.save('Universe-GPT-tf.keras', save_format='tf')