import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
import torch.nn.functional as F


def generate_linear_equations(number_of_samples, max_int=100):
    a_values = np.random.randint(1, max_int, size=(number_of_samples, 1))
    x_values = np.random.randint(1, max_int, size=(number_of_samples, 1))
    b_values = np.random.randint(1, max_int, size=(number_of_samples, 1))

    c_values = a_values * x_values + b_values  # Compute c in ax + b = c
    
    equations = np.concatenate((a_values, b_values, c_values), axis=1)
    return equations, x_values

# 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)
        h = h[:, -1, :]  # taking the last sequence output
        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):
        # Flatten if needed (in case input comes from LSTM with sequence length)
        if x.dim() > 2:
            x = x.view(x.size(0), -1)
        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):
        # Compute the gating weights
        gating_scores = self.gating_network(x).unsqueeze(-1) # [batch_size, num_experts, 1]
        # Compute the output from each expert for every example
        expert_outputs = [expert(x) for expert in self.experts]
        expert_outputs = torch.stack(expert_outputs, dim=-1)  # [batch_size, output_size, num_experts]
        
        # Multiply gating scores by expert outputs and sum the result
        # torch.bmm expects [batch_size, n, m] and [batch_size, m, p]
        # here it will do [batch_size, output_size, num_experts] bmm [batch_size, num_experts, 1]
        mixed_output = torch.bmm(expert_outputs, gating_scores)

        # Now mixed_output has shape [batch_size, output_size, 1], we should squeeze it to [batch_size, output_size]
        mixed_output = mixed_output.squeeze(-1)

        return mixed_output
    
class SingleLSTM(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers=3):  # This model will have more parameters
        super(SingleLSTM, self).__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        h, _ = self.lstm(x)
        h = h[:, -1, :]  # taking the last sequence output
        return self.fc(h)
    
# Simple Positional Encoding (not learnable, for simplicity)
def positional_encoding(seq_len, d_model, device):
    pos = torch.arange(seq_len, dtype=torch.float, device=device).reshape(1, -1, 1)
    dim = torch.arange(d_model, dtype=torch.float, device=device).reshape(1, 1, -1)
    phase = pos / (1e4 ** (dim // d_model))

    return torch.where(dim.long() % 2 == 0, torch.sin(phase), torch.cos(phase))

class SimpleTransformer(nn.Module):
    def __init__(self, input_size, d_model, output_size, nhead, dim_feedforward, num_encoder_layers=1):
        super(SimpleTransformer, self).__init__()
        self.d_model = d_model
        self.input_fc = nn.Linear(input_size, d_model)  # Linear layer to project to model dimension
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            batch_first=True)  # Added batch_first=True
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
        self.output_fc = nn.Linear(d_model, output_size)

    def forward(self, x):
        seq_len, batch_size, _ = x.size()
        x = self.input_fc(x) + positional_encoding(seq_len, self.d_model, x.device)

        # Adding positional encodings
        pos_enc = positional_encoding(seq_len, self.d_model, x.device)
        x = x + pos_enc
        
        # Passing the input through the Transformer
        transformer_output = self.transformer_encoder(x)

        # Taking only the result from the last token in the sequence
        transformer_output = transformer_output[-1, :, :]

        # Final fully connected layer to produce the output
        output = self.output_fc(transformer_output)
        return output.squeeze(-1) # Remove the last dimension for compatibility with the target
    
class TransformerExpert(nn.Module):
    def __init__(self, input_size, d_model, output_size, nhead, dim_feedforward, num_encoder_layers=1):
        super(TransformerExpert, self).__init__()
        self.d_model = d_model
        self.input_fc = nn.Linear(input_size, d_model)  # Linear layer to project to model dimension
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, batch_first=True)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
        self.output_fc = nn.Linear(d_model, output_size)

    def forward(self, x):
        x = self.input_fc(x)
        x += positional_encoding(x.size(1), self.d_model, x.device)
        transformer_output = self.transformer_encoder(x)
        transformer_output = transformer_output[:, -1, :]
        output = self.output_fc(transformer_output)
        return output

class MoETransformer(nn.Module):
    def __init__(self, input_size, d_model, output_size, nhead, dim_feedforward, num_experts, num_encoder_layers=1):
        super(MoETransformer, self).__init__()
        self.num_experts = num_experts
        self.experts = nn.ModuleList([TransformerExpert(input_size, d_model, output_size, nhead, dim_feedforward, num_encoder_layers)
                                      for _ in range(num_experts)])
        self.gating_network = GatingNetwork(input_size, num_experts)

    def forward(self, x):
        # Compute the gating weights
        gating_scores = self.gating_network(x).unsqueeze(-1)  # [batch_size, num_experts, 1]
        
        # Compute the output from each expert for every example
        expert_outputs = [expert(x).unsqueeze(2) for expert in self.experts]  # [batch_size, output_size, 1]
        expert_outputs = torch.cat(expert_outputs, dim=2)   # [batch_size, output_size, num_experts]
        
        # Multiply gating scores by expert outputs and sum the result
        mixed_output = torch.bmm(expert_outputs, gating_scores)  # [batch_size, output_size, 1]
        
        # Now mixed_output has shape [batch_size, output_size, 1], we squeeze it to [batch_size, output_size]
        mixed_output = mixed_output.squeeze(-1)
        return mixed_output
    
def train_model(model, criterion, optimizer, num_epochs, batch_size, equations_tensor, solutions_tensor):
    model.train()
    training_losses = []
    for epoch in range(num_epochs):
        total_loss = 0

        # Shuffle the data
        perm = torch.randperm(equations_tensor.size(0))
        equations_shuffled = equations_tensor[perm]
        solutions_shuffled = solutions_tensor[perm]

        for i in range(0, equations_tensor.size(0), batch_size):
            # Slice out the batches
            x_batch = equations_shuffled[i:i+batch_size]
            y_batch = solutions_shuffled[i:i+batch_size]

            optimizer.zero_grad()
            predictions = model(x_batch)
            predictions = predictions.squeeze()  # Ensure predictions and targets have the same size
            loss = criterion(predictions, y_batch.view(-1))  # Flatten the target tensor
            loss.backward()
            optimizer.step()

            total_loss += loss.item()
        
        average_loss = total_loss / (equations_tensor.size(0) / batch_size)
        training_losses.append(average_loss)
        print(f"Epoch {epoch+1}, Loss: {average_loss}")
        
    return training_losses

# Parameters for the Transformer model
d_model = 128
nhead = 4
dim_feedforward = 256

input_size = 3  # a, b, and c from equation ax + b = c
hidden_size = 32
output_size = 1  # we are predicting the value of x
num_experts = 10

# Training parameters
num_epochs = 20
batch_size = 128

# Reshape the equations to (batch_size, seq_length, features)
# For simplicity, we consider each equation as a sequence of length 1
# Generate the data only once before training
equations_tensor, solutions_tensor = generate_linear_equations(number_of_samples=5000)

# Convert numpy arrays to torch tensors
equations_tensor = torch.from_numpy(equations_tensor).float()
solutions_tensor = torch.from_numpy(solutions_tensor).float()

# Reshape the equations and solutions to (batch_size, seq_length, features)
train_equations_tensor = equations_tensor.view(-1, 1, input_size)  # reshaping once, correctly
train_solutions_tensor = solutions_tensor.view(-1, 1)              # reshaping for solutions to match batch size

# Instantiate the MoE-Transformer model
moe_transformer_model = MoETransformer(input_size=input_size, d_model=d_model, output_size=output_size, nhead=nhead, dim_feedforward=dim_feedforward, num_experts=num_experts)
moe_transformer_optimizer = torch.optim.Adam(moe_transformer_model.parameters())
moe_transformer_criterion = nn.MSELoss()

# Instantiate the SimpleTransformer model
transformer_model = SimpleTransformer(input_size=input_size, d_model=d_model, output_size=output_size, nhead=nhead, dim_feedforward=dim_feedforward)
transformer_optimizer = torch.optim.Adam(transformer_model.parameters())
transformer_criterion = nn.MSELoss()

# Instantiate the MoE model
moe_model = MixtureOfExperts(input_size, hidden_size*2, output_size, num_experts)
moe_optimizer = torch.optim.Adam(moe_model.parameters())
moe_criterion = nn.MSELoss()

# Instantiate the Single General LSTM model
large_model = SingleLSTM(input_size, hidden_size*2, output_size, num_layers=4)  # Adjusted to have comparable number of parameters
large_optimizer = torch.optim.Adam(large_model.parameters())
large_criterion = nn.MSELoss()

# Print the parameter count of the MoE-Transformer
print("MoE-Transformer Parameter Count:", sum(p.numel() for p in moe_transformer_model.parameters()))

# Print the parameter count of the Transformer
print("Transformer Parameter Count:", sum(p.numel() for p in transformer_model.parameters()))

# Print the parameter count of the MoE LSTM
print("MoE LSTM Parameter Count:", sum(p.numel() for p in moe_model.parameters()))

# Print the parameter count of the Single LSTM
print("Single LSTM Parameter Count:", sum(p.numel() for p in large_model.parameters()))

# Train the MoE Transformer model
moe_transformer_losses = train_model(moe_transformer_model, moe_transformer_criterion, moe_transformer_optimizer,num_epochs, batch_size, train_equations_tensor, train_solutions_tensor)

# Train the Transformer model
transformer_losses = train_model(transformer_model, transformer_criterion, transformer_optimizer, num_epochs, batch_size, train_equations_tensor, train_solutions_tensor)

# Train the MoE model
moe_losses = train_model(moe_model, moe_criterion, moe_optimizer, num_epochs, batch_size, train_equations_tensor, train_solutions_tensor)

# Train the Single LSTM model
large_losses = train_model(large_model, large_criterion, large_optimizer, num_epochs, batch_size, train_equations_tensor, train_solutions_tensor)

# Plot the results, as before
plt.figure(figsize=(12, 8))
plt.plot(moe_losses, label="Mixture of Experts Loss")
plt.plot(large_losses, label="Single LSTM Model Loss")
plt.plot(transformer_losses, label="Simple Transformer Model Loss")
plt.plot(moe_transformer_losses, label="MoE Transformer Model Loss")
plt.xlabel("Epoch")
plt.ylabel("MSE Loss")
plt.legend()
plt.title("Model Loss Comparison")
plt.show()