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>"]]
    vocab_size = len(vocab)
    print("Vocabulary size:", vocab_size)
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 = 4
state_dim = 128
hidden_dim = 128
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
            
            for i in range(num_batches):
                batch_start_index = i * batch_size
                inputs = data[batch_start_index:batch_start_index + batch_size, :-1]  # Exclude last column for inputs
                targets = data[batch_start_index:batch_start_index + batch_size, 1:]  # Exclude first column for targets
                targets = targets.reshape(-1)  # Flatten the targets to [batch_size * seq_length]

                optimizer.zero_grad()
                outputs = model(inputs)
                outputs = outputs.reshape(-1, output_features)  # [batch_size * seq_length, output_features]
                loss = loss_fn(outputs, targets)

                loss.backward()
                optimizer.step()

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

        return loss_history

    output_features = vocab_size  # Make sure output_features matches the number of classes for prediction
    epochs = 10
    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)