import torch
import torch.nn as nn
import torch.nn.functional as F
from torchtext.data.utils import get_tokenizer
from torchtext.vocab import build_vocab_from_iterator
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm

# Define the Expert module
class Expert(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super(Expert, self).__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.activation = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, input_dim)

    def forward(self, x):
        x = self.activation(self.fc1(x))
        x = self.fc2(x)
        return x

# Define the MoELayer module
class MoELayer(nn.Module):
    def __init__(self, num_experts, input_dim, hidden_dim):
        super(MoELayer, self).__init__()
        self.experts = nn.ModuleList([Expert(input_dim, hidden_dim) for _ in range(num_experts)])
        self.gate = nn.Linear(input_dim, num_experts)

    def forward(self, x):
        gating_distribution = F.softmax(self.gate(x), dim=1)
        expert_outputs = torch.stack([expert(x) for expert in self.experts], dim=1)
        output = torch.einsum('bi,bij->bj', gating_distribution, expert_outputs)  # This efficiently computes the weighted sum of expert outputs
        return output

# Define the SelectionMechanism module
class SelectionMechanism(nn.Module):
    def __init__(self, input_dim, state_dim):
        super(SelectionMechanism, self).__init__()
        self.selection_layer = nn.Linear(input_dim + state_dim, state_dim)

    def forward(self, state, u):
        # Ensure both state and u have a batch dimension (state already has it)
        combined = torch.cat((state, u), dim=1)
        new_state = self.selection_layer(combined)
        return new_state

# Define the StateSpaceMamba module
class StateSpaceMamba(nn.Module):
    def __init__(self, input_dim, state_dim, num_experts, hidden_dim):
        super(StateSpaceMamba, self).__init__()
        self.state_dim = state_dim
        self.state = nn.Parameter(torch.zeros(1, state_dim))
        self.input_layer = nn.Linear(input_dim, state_dim)
        self.selection_mechanism = SelectionMechanism(input_dim, state_dim)
        self.moe_layer = MoELayer(num_experts, state_dim, hidden_dim)

    def forward(self, input_sequence):
        batch_size = input_sequence.size(0)  # Assuming the input is [batch_size, seq_length, input_dim]
        outputs = []
        state = torch.zeros(batch_size, self.state_dim).to(input_sequence.device)  # Initialize state for the current batch

        # Debugging: print shapes before loop
        print("input_sequence shape at forward entry:", input_sequence.shape)
        print("Initial state shape at forward entry:", state.shape)

        for u in input_sequence.split(1, dim=1):  # Corrected: split along dimension 1 (time/sequence)
            u = u.squeeze(1)  # Corrected: Remove the sequence dimension which is now dim 1 due to split
            # Debugging: Print shapes inside loop
            print("Input shape (u) at loop:", u.shape)
            state = self.selection_mechanism(state, u)
            print("State shape after selection mechanism:", state.shape)
            u_projected = self.input_layer(u)
            print("Input projection (u_projected) shape:", u_projected.shape)
            output = self.moe_layer(u_projected + state)
            print("Output shape after MoE layer:", output.shape)
            
            outputs.append(output)

        # Debugging: print shape of individual outputs and final stacked output
        print("Individual outputs shapes:", [o.shape for o in outputs])
        output_tensor = torch.stack(outputs, dim=1)  # Corrected: stack along dimension 1 (time/sequence)
        print("Stacked output shape:", output_tensor.shape)

        return output_tensor

# Load the dataset
try:
    print("Loading dataset...")
    with open("C:/Users/Dell-PC/Desktop/The-Physics-Dataset.txt", "r") as f:
        text = f.read()
except FileNotFoundError:
    print("Error: Dataset file not found.")
    exit(1)

# Tokenize the text
try:
    print("Tokenizing text...")
    tokenizer = get_tokenizer("basic_english")
    tokens = tokenizer(text.lower())
except Exception as e:
    print("Error occurred during tokenization:", str(e))
    exit(1)

# Build the vocabulary
try:
    print("Building vocabulary...")
    vocab = build_vocab_from_iterator([tokens], specials=["<unk>", "<pad>", "<sos>", "<eos>"])
    vocab.set_default_index(vocab["<unk>"])
    numericalized_tokens = [vocab["<sos>"]] + [vocab[token] for token in tokens] + [vocab["<eos>"]]
except Exception as e:
    print("Error occurred during vocabulary building:", str(e))
    exit(1)

# Prepare the input sequence
batch_size = 24
seq_length = 12  # This is the “token” size we observed in the debug outputs
vocab_size = len(vocab)
num_batches = (len(numericalized_tokens) - 1) // (batch_size * seq_length)

if len(numericalized_tokens) < (batch_size * seq_length):
    print(f"Error: Dataset is too small for the chosen batch size and sequence length. Tokens available: {len(numericalized_tokens)}")
    exit(1)

# Calculate the number of batches and ensure it’s not zero
if num_batches == 0:
    print(f"Error: Number of batches is zero. Check if your dataset is large enough or consider reducing batch size/sequence length.")
    exit(1)

numericalized_tokens = numericalized_tokens[:num_batches * batch_size * seq_length]
input_sequence = torch.tensor(numericalized_tokens).view(-1, batch_size, seq_length)
input_sequence = input_sequence.float()  # Ensure input_sequence is of type float

print("input_sequence shape:", input_sequence.shape)
print("input_sequence dimension:", input_sequence.dim())

# Define the model using actual input_dim which corresponds to seq_length
num_experts = 5
state_dim = 256
hidden_dim = 256
input_dim = seq_length  # Corrected to use seq_length instead of vocab_size
model = StateSpaceMamba(input_dim, state_dim, num_experts, hidden_dim)
learning_rate = 0.001

print("model state shape:", model.state.shape)
print("model state dimension:", model.state.dim())

# Training loop
try:
    print("Training model…")
    def train(model, data, num_batches, batch_size, seq_length, output_features, epochs):
        model.train()
        loss_fn = nn.CrossEntropyLoss()
        optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
        loss_history = []

        for epoch in tqdm(range(epochs)):
            total_loss = 0.0
            # Debugging: print data tensor shape
            print("Data tensor shape:", data.shape)
            
            for i in range(0, data.size(1) - seq_length, seq_length):
                # Debugging: print current index and input-target shapes
                print(f"Current index: {i}")
                inputs = data[:, i:i+seq_length]
                print("Inputs shape:", inputs.shape)
                targets = data[:, i+1:i+1+seq_length]
                targets = targets[:, 1:] # remove the first element of each sequence
                targets = targets.reshape(-1) # flatten the tensor
                targets = targets.long() # convert to long tensor


                optimizer.zero_grad()
                outputs = model(inputs)
                # No need for transpose since the sequence is already the second dimension out of model.
                # We reshape outputs to match the total number of predictions for all sequences in all batches.
                outputs = outputs.reshape(-1, output_features)  # Flatten the ouput predictions
                
                loss = loss_fn(outputs, targets)
                loss.backward()
                optimizer.step()
                # Validate that vocab_size actually matches the last dimension of outputs
                assert outputs.size(-1) == vocab_size, "Mismatch between output size and vocab size."

                total_loss += loss.item()
            epoch_loss = total_loss / num_batches
            loss_history.append(epoch_loss)
            print(f"Epoch {epoch}, Loss: {epoch_loss}")

        return loss_history

    # Same number of features as the last layer of the model, typically the logits layer.
    output_features = 256
    epochs = 10
    # Correct the train function call by replacing vocab_size with output_features
    loss_history = train(model, input_sequence, num_batches, batch_size, seq_length, output_features, epochs)

    # Plot loss history
    plt.plot(loss_history)
    plt.xlabel("Epochs")
    plt.ylabel("Loss")
    plt.title("Training Loss History")
    plt.show()

except Exception as e:
    print("Error occurred during training:", str(e))
    exit(1)

# Generate function
try:
    print("Generating text...")
    def generate(model, start_seq, length, temperature=0.8):
        model.eval()
        state = model.state
        numericalized_start = torch.tensor([vocab[token] for token in tokenizer(start_seq.lower())], dtype=torch.long).unsqueeze(0)
        input_seq = numericalized_start

        generated_sequence = []
        for i in range(length):
            output = model(input_seq)

            # Sample from the network as a multinomial distribution
            output_prob = F.softmax(output[:, -1] / temperature, dim=-1).detach()
            predicted_index = torch.multinomial(output_prob, 1).item()

            generated_sequence.append(predicted_index)

            input_seq = torch.cat((input_seq, torch.tensor([[predicted_index]])), dim=1)
        return " ".join(vocab.lookup_tokens(generated_sequence))

    # Generating a sample
    start_seq = "The fundamental concepts of physics"  # Start of sequence to begin text generation
    generated_text = generate(model, start_seq, length=100)  # Length of generated sequence
    print(generated_text)

except Exception as e:
    print("Error occurred during text generation:", str(e))
    exit(1)