import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader, Dataset
from collections import Counter
import json




# Expert LSTM Model
class LSTMExpert(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(LSTMExpert, self).__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        h, _ = self.lstm(x)
        # Shape of h: [batch_size, sequence_length, hidden_size]
        # No longer taking the last sequence output, transform for all time steps
        h = h.reshape(-1, h.shape[2])  # Flatten to two dimensions: [batch_size * sequence_length, hidden_size]
        return self.fc(h)

# Gating Network
class GatingNetwork(nn.Module):
    def __init__(self, input_size, num_experts):
        super(GatingNetwork, self).__init__()
        self.fc = nn.Linear(input_size, num_experts)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        if x.dim() > 2:
            x = x.view(x.size(0), -1)  # Flatten the tensor correctly
        x = self.fc(x)
        return self.softmax(x)

# Mixture of Experts Model
class MixtureOfExperts(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_experts):
        super(MixtureOfExperts, self).__init__()
        self.num_experts = num_experts
        self.experts = nn.ModuleList([LSTMExpert(input_size, hidden_size, output_size) for _ in range(num_experts)])
        self.gating_network = GatingNetwork(input_size, num_experts)

    def forward(self, x):
        # Flatten x for gating input: [batch_size * sequence_length, input_size]
        gating_input = x.contiguous().view(-1, x.shape[-1])
        # Get gating scores for each expert: [batch_size * sequence_length, num_experts, 1]
        gating_scores = self.gating_network(gating_input).view(-1, self.num_experts, 1)

        # Compute outputs from each expert
        expert_outputs = [expert(x) for expert in self.experts]  # List of [batch_size * sequence_length, vocab_size] tensors
        # Stack along a new dimension for experts: [batch_size * sequence_length, vocab_size, num_experts]
        expert_outputs = torch.stack(expert_outputs, dim=2)

        # Calculate weighted sum of expert outputs: [batch_size * sequence_length, vocab_size]
        mixed_output = torch.bmm(expert_outputs, gating_scores)  # [batch_size * sequence_length, vocab_size, 1]
        mixed_output = mixed_output.squeeze(-1)  # Remove last singleton dimension

        return mixed_output


class QAJsonlDataset(Dataset):
    def __init__(self, path, seq_len):
        self.seq_len = seq_len
        self.pairs = self.load_data(path)
        self.vocab, self.idx2token = self.build_vocab([word for pair in self.pairs for word in pair])
        self.tokenized_pairs = [(self.tokenize(q), self.tokenize(a)) for q, a in self.pairs]

    def load_data(self, path):
        pairs = []
        with open(path, "r", encoding="utf-8") as f:
            for line in f:
                data = json.loads(line.strip())
                question, answer = data.get("question", ""), data.get("answer", "")
                pairs.append((question.split(), answer.split()))
        return pairs

    def tokenize(self, words):
        # Tokenize a sentence and pad if necessary
        tokens = [self.vocab.get(w, self.vocab["<unk>"]) for w in words]  # Use <unk> for unknown words
        if len(tokens) < self.seq_len:
            tokens += [self.vocab["<pad>"]] * (self.seq_len - len(tokens))
        else:
            tokens = tokens[:self.seq_len-1] + [self.vocab["<eos>"]]  # Assuming you have an <eos> token in your vocab
        return tokens

    def build_vocab(self, words):
        # You might include “<unk>”, “<pad>” etc. in the vocabulary beforehand
        counts = Counter(words)
        vocab = {word: i for i, (word, _) in enumerate(counts.most_common(), 2)}
        vocab["<unk>"] = 0
        vocab["<pad>"] = 1
        idx2token = {i: word for word, i in vocab.items()}
        return vocab, idx2token
        
    def __len__(self):
        return len(self.tokenized_pairs)

    def __getitem__(self, idx):
        tokenized_question, tokenized_answer = self.tokenized_pairs[idx]
        return torch.tensor(tokenized_question, dtype=torch.long), torch.tensor(tokenized_answer, dtype=torch.long)

def collate_fn(batch):
    questions, answers = zip(*batch)
    questions = pad_sequence(questions, batch_first=True, padding_value=0)
    answers = pad_sequence(answers, batch_first=True, padding_value=0)
    return questions, answers

# Set the path to your text file and define sequence length
path_to_text = 'GSM2K.jsonl'  # replace with the path to your text file
seq_len = 32  # sequence length

# Create a dataset and data loader
dataset = QAJsonlDataset(path_to_text, seq_len)
data_loader = DataLoader(dataset, batch_size=32, shuffle=True, collate_fn=collate_fn)
    
def train_model_TXT(model, criterion, optimizer, num_epochs, data_loader):
    model.train()
    for epoch in range(num_epochs):
        total_loss = 0

        for i, (inputs, targets) in enumerate(data_loader):
            optimizer.zero_grad()

            # Forward pass
            predictions = model(inputs)

            # Flatten targets to be [batch_size * sequence_length]
            targets = targets.view(-1)

            # Flatten predictions to be [batch_size * sequence_length, num_classes]
            predictions = predictions.view(-1, predictions.size(-1))

            # Loss computation
            loss = criterion(predictions, targets)
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        average_loss = total_loss / len(data_loader.dataset)
        print(f"Epoch {epoch+1}, Loss: {average_loss}")

def generate_text(model, dataset, seed_text, num_generate, temperature=1.0):
    model.eval()  # Put the model in evaluation mode
    
    # List to store the generated tokens
    generated_tokens = []

    # Initial sequence (prefix) to start the generation process
    input_sequence = [dataset.vocab.get(word, dataset.vocab["<pad>"]) for word in seed_text.split()]  # Convert to token IDs
    current_sequence = torch.tensor(input_sequence, dtype=torch.long).unsqueeze(0)

    # Generate num_generate tokens
    for _ in range(num_generate):
        # Forward pass through the model
        with torch.no_grad():
            output = model(current_sequence)
        
        # Get probabilities, apply temperature scaling, and sample from the distribution
        probabilities = F.softmax(output[-1] / temperature, dim=0).detach()
        next_token_idx = torch.multinomial(probabilities, 1).item()

        # Append token to the current sequence and to the generated tokens
        generated_tokens.append(next_token_idx)
        current_sequence = torch.cat((current_sequence, torch.tensor([[next_token_idx]])), 1)
    
    # Convert tokens to words
    generated_text = " ".join([dataset.idx2token.get(token, "<unk>") for token in generated_tokens])  # Use .get() to provide a default value for missing keys
    return generated_text

# Integration of MoE with the rest of the code
vocab_size = len(dataset.vocab)  # The size of the vocabulary
embedding_dim = 128  # Size of embeddings
hidden_size = 256  # Number of features in the hidden state of the LSTM
num_experts = 2  # Define the number of LSTM experts you want in your MoE model

# Instantiate the MixtureOfExperts model
mixture_of_experts_model = MixtureOfExperts(embedding_dim, hidden_size, vocab_size, num_experts)

# Wrap the experts and the embedding layer into a new combined model, replacing SingleLSTM_TXT
class MoEModel(nn.Module):
    def __init__(self, embedding_dim, vocab_size, MoE):
        super(MoEModel, self).__init__()
        self.embedding = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embedding_dim)
        self.MoE = MoE

    def forward(self, x):
        embedded = self.embedding(x)
        output = self.MoE(embedded)
        return output

# Instantiate the MoEModel with embeddings and mixture of experts
moe_model = MoEModel(embedding_dim, vocab_size, mixture_of_experts_model)

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

# Example usage with your model:
total_params = count_parameters(moe_model)
print(f"Total trainable parameters: {total_params}")

# Training parameters
num_epochs = 2
batch_size = 32
learning_rate = 1e-3

# Define Loss Function and Optimizer for MoE model
criterion = nn.CrossEntropyLoss()  # Use CrossEntropyLoss for classification tasks
optimizer = torch.optim.Adam(moe_model.parameters(), lr=learning_rate)

# Replace references to single_lstm_txt_model with moe_model
# Train the model with the text data
train_model_TXT(moe_model, criterion, optimizer, num_epochs, data_loader)

# Start a loop for the interactive chat-like text generation
while True:
    try:
        # Get user input
        seed_text = input("Enter seed text (type 'quit' to stop): ")
        
        # Check if user wants to quit the interaction
        if seed_text.lower() == "quit":
            print("Exiting text generation chat.")
            break
        
        # User input is not empty and not “quit”, generate text
        if seed_text.strip():
            num_generate = 16  # Number of words to generate
            temperature = 0.5  # Sampling temperature, higher will increase diversity
            
            # Use the trained model to generate text
            generated_text = generate_text(moe_model, dataset, seed_text, num_generate, temperature)
            print("Generated Text:", generated_text)
        else:
            print("Seed text cannot be empty.")
    except KeyboardInterrupt:
        # Handle KeyboardInterrupt (Ctrl+C) to gracefully exit
        print("\nExiting text generation chat.")
        break