import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from tqdm import tqdm
import matplotlib.pyplot as plt
import string
import json
import os

# Define parameters and directory paths
DATASET_CHUNKS = 3805  # New hyperparameter to control dataset size (number of chunks)
SEQUENCE_LENGTH = 60
BATCH_SIZE = 32
HIDDEN_SIZE = 768
NUM_LAYERS = 2
LEARNING_RATE = 0.001
NUM_EPOCHS = 2
GRAD_CLIP = 5.0  # Gradient clipping to prevent exploding gradients
CHECKPOINT_FOLDER = "checkpoints"
os.makedirs(CHECKPOINT_FOLDER, exist_ok=True)
model_name = "Talk-model-1"

# Replace ‘your_file.txt’ with the path to your actual text file.
file_path = "Long-Run-Talk-Cutdown.txt"

# Print the start of data preparation:
print("\n--- Starting data preparation ---")

# Read the file and filter out blank lines
with open(file_path, "r", encoding="utf-8") as file:
    lines = file.readlines()
    text = "".join(line for line in lines if not line.isspace())

# Convert text to ASCII values
ascii_characters = [ord(c) for c in text if c in string.printable[:-6]] # Make sure to include only printable characters
print(f"Converted text to ASCII characters: {len(ascii_characters)} characters")

# Define the total number of sequences based on the sequence length
total_num_sequences = len(ascii_characters) - SEQUENCE_LENGTH

# Calculate the number of sequences to use per chunk
sequences_per_chunk = total_num_sequences // DATASET_CHUNKS
# This will be the actual first dimension for the input and target tensors
usable_sequences = sequences_per_chunk * DATASET_CHUNKS

# Prepare the datasets
inputs = []
targets = []
for start_idx in range(0, usable_sequences, sequences_per_chunk):
    for i in range(start_idx, start_idx + sequences_per_chunk):
        input_seq = ascii_characters[i:i+SEQUENCE_LENGTH]
        target_seq = ascii_characters[i+1:i+SEQUENCE_LENGTH+1]
        inputs.append(torch.tensor(input_seq, dtype=torch.long))
        targets.append(torch.tensor(target_seq, dtype=torch.long))

input_tensor = torch.stack(inputs)
target_tensor = torch.stack(targets)
print(f"Created input tensor: {input_tensor.shape}")
print(f"Created target tensor: {target_tensor.shape}")

# Load data into DataLoader for batch processing
dataset = TensorDataset(input_tensor, target_tensor)
data_loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
print("\n--- Data preparation complete ---")

# Define the sophisticated RNN model architecture
class AdvancedRNN(nn.Module):
    def __init__(self, vocab_size, hidden_size, num_layers):
        super(AdvancedRNN, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.embedding = nn.Embedding(vocab_size, hidden_size)
        self.rnn = nn.LSTM(hidden_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, vocab_size)

    def forward(self, x, hidden):
        x = self.embedding(x)
        out, hidden = self.rnn(x, hidden)
        out = self.fc(out.reshape(out.size(0) * out.size(1), out.size(2)))
        return out, hidden

    def init_hidden(self, batch_size):
        weight = next(self.parameters())
        hidden = (weight.new_zeros(self.num_layers, batch_size, self.hidden_size),
                  weight.new_zeros(self.num_layers, batch_size, self.hidden_size))
        return hidden

# Generate vocabulary with printable ASCII characters
vocab_chars = string.printable[:-6]  # Exclude the last 6 non-printable characters
VOCAB_SIZE = len(vocab_chars)  # Set the size of the vocabulary based on actual characters used

# Map each character to a unique index
char_to_index = {char: idx for idx, char in enumerate(vocab_chars)}

# Initialize model with correct VOCAB_SIZE
model = AdvancedRNN(VOCAB_SIZE, HIDDEN_SIZE, NUM_LAYERS)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)

# Calculate and print the total number of trainable parameters
total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Total trainable parameters: {total_params}")

print("\n--- Starting training loop ---")
# Training loop with added production-ready features and tqdm progress
def train_model(model, data_loader, optimizer, loss_fn, epochs, batch_size, model_name):
    loss_history = [] # Initialize a list to store average epoch losses
    for epoch in range(epochs):
        model.train()
        hidden = model.init_hidden(batch_size)
        running_loss = 0.0
        # Wrap the data_loader with tqdm for a progress bar
        progress_bar = tqdm(data_loader, desc=f"Epoch {epoch + 1}/{epochs}")

        for input_batch, target_batch in progress_bar:
            if input_batch.size(0) != batch_size:
                hidden = model.init_hidden(input_batch.size(0))

            optimizer.zero_grad()
            output, hidden = model(input_batch, hidden)
            hidden = tuple([h.detach() for h in hidden]) # Detach hidden state
            loss = loss_fn(output, target_batch.view(-1))
            loss.backward()

            # Clip gradients to mitigate exploding gradients
            nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
            optimizer.step()

            running_loss += loss.item()
            progress_bar.set_postfix(loss=loss.item())

        # Calculate and store the average epoch loss
        epoch_loss = running_loss / len(data_loader)
        loss_history.append(epoch_loss)
        print(f"Epoch {epoch + 1}/{epochs}, Loss: {epoch_loss:.4f}")

    # Save the last model checkpoint instead of every epoch
    checkpoint_path = os.path.join(CHECKPOINT_FOLDER, f"{model_name}.pth")
    torch.save(model.state_dict(), checkpoint_path)

    # Save model’s configuration in a JSON file
    config_path = os.path.join(CHECKPOINT_FOLDER, f"{model_name}_config.json")
    model_config = {
        "vocab_size": VOCAB_SIZE,
        "hidden_size": HIDDEN_SIZE,
        "num_layers": NUM_LAYERS,
        "batch_size":BATCH_SIZE,
        "Total_params":total_params
    }
    with open(config_path, "w") as json_file:
        json.dump(model_config, json_file)

    # Generate vocabulary with printable ASCII characters
    vocab_chars = string.printable[:-6]  # Exclude the last 6 non-printable characters
    VOCAB_SIZE = len(vocab_chars)  # Update the size of the vocabulary

    # Save vocabulary in a TXT file
    vocab_path = os.path.join(CHECKPOINT_FOLDER, f"{model_name}_vocab.txt")
    with open(vocab_path, "w", encoding="utf-8") as vocab_file:
        vocab_file.write(vocab_chars)

    # Now, build a character to index mapping
    char_to_index = {char: idx for idx, char in enumerate(vocab_chars)}

    # Plot training loss per epoch outside the training loop
    plt.figure() # Create a new figure
    plt.plot(range(1, epochs + 1), loss_history, label="Training Loss")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.legend()
    plt.grid(True)
    plt.title("Training Loss Over Epochs")
    plt.show()

    print("\n--- Training loop complete ---")

    return model  # Return the trained model

# Predict the next character
def predict(model, char, hidden=None, temperature=1.0):
    # Convert to ASCII and then tensor
    input_tensor = torch.tensor([ord(char)], dtype=torch.long).unsqueeze(0)
    if hidden is None:
        hidden = model.init_hidden(1)
    output, hidden = model(input_tensor, hidden)
    output_dist = output.view(-1).div(temperature).exp()
    top_i = torch.multinomial(output_dist, 1)[0]
    predicted_char = chr(top_i)
    return predicted_char, hidden

# Generate text
def generate_text(model, start_char="P", length=200):
    model.eval()
    hidden = model.init_hidden(1)
    input_char = start_char
    for i in range(length):  
        input_char, hidden = predict(model, input_char, hidden)
        print(input_char, end="")

if __name__ == "__main__":
    # Train the model
    trained_model = train_model(model, data_loader, optimizer, loss_fn, NUM_EPOCHS, BATCH_SIZE, model_name)

    # Generate text from the trained model
    generate_text(trained_model, start_char="P", length=200)

    # Save vocabulary in a TXT file (this only needs to be done once)
    vocab_path = os.path.join(CHECKPOINT_FOLDER, f"{model_name}_vocab.txt")
    with open(vocab_path, "w", encoding="utf-8") as vocab_file:
        vocab_file.write(vocab_chars)