import torch, os
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np

# Configuration parameters
NUM_LAYERS = 12
D_MODEL = 512
NUM_HEADS = 8
DFF = 2048  # Feed-forward network dimension
VOCAB_SIZE = 50257
SEQ_LENGTH = 128

# Define a minimal transformer model
class TransformerModel(nn.Module):
    def __init__(self, vocab_size, d_model, num_heads, num_layers, dff):
        super(TransformerModel, self).__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.transformer_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=num_heads, dim_feedforward=dff)
        self.transformer = nn.TransformerEncoder(self.transformer_layer, num_layers=num_layers)
        self.fc_out = nn.Linear(d_model, vocab_size)

    def forward(self, src):
        x = self.embedding(src) * (self.embedding.embedding_dim ** 0.5)
        x = self.transformer(x)
        x = self.fc_out(x)
        return x

# A generator to simulate data streaming
class TextDatasetIter:
    def __init__(self, file_path):
        self.file_path = file_path
        with open(file_path, 'r', encoding='utf-8') as f:
            self.text = f.read()
        self.char2idx = {u: i for i, u in enumerate(set(self.text))}
        self.data = [self.char2idx[c] for c in self.text]

    def __len__(self):
        return len(self.data) - SEQ_LENGTH

    def __iter__(self):
        for i in range(len(self.data) - SEQ_LENGTH):
            yield (
                torch.tensor(self.data[i:i + SEQ_LENGTH], dtype=torch.long),
                torch.tensor(self.data[i + 1:i + SEQ_LENGTH + 1], dtype=torch.long)
            )

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = TransformerModel(VOCAB_SIZE, D_MODEL, NUM_HEADS, NUM_LAYERS, DFF).to(device)

optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

dataset_iterator = TextDatasetIter("training.txt")
batch_size = 16

# Training loop
def train_model(model, dataset_iterator, optimizer, criterion, num_epochs=10, save_every=1):
    model.train()
    for epoch in range(num_epochs):
        batch_loss = 0
        current_batch = []

        for idx, (input_seq, target_seq) in enumerate(dataset_iterator):
            current_batch.append((input_seq, target_seq))

            if len(current_batch) == batch_size:  # Process batch
                batch_inputs, batch_targets = zip(*current_batch)
                batch_inputs = torch.stack(batch_inputs).to(device)
                batch_targets = torch.stack(batch_targets).to(device)

                optimizer.zero_grad()

                output = model(batch_inputs.permute(1, 0))
                loss = criterion(output.view(-1, VOCAB_SIZE), batch_targets.view(-1))
                loss.backward()
                optimizer.step()

                batch_loss += loss.item()
                current_batch = []  # Reset batch

                if idx % save_every == 0:
                    print(f"[Epoch {epoch+1}, Step {idx}] Loss: {loss.item():.4f}")

        print(f"Epoch {epoch+1} Complete with Average Loss: {batch_loss / len(dataset_iterator):.4f}")

train_model(model, dataset_iterator, optimizer, criterion)


alter this to have better training thing with like a estimation, iterations per second, etc. and make it save the model checkpoint and autodelete older ones, maximum amount of checkpoints is 3 and it saves a new checkpoint once every 1000 iters