import jax
import jax.numpy as jnp
import jax.random as random
from flax import linen as nn
from flax.training import train_state
import optax
import matplotlib.pyplot as plt
from tqdm import tqdm
import numpy as np

# Define standard transformer components
class MultiHeadSelfAttention(nn.Module):
    d_model: int
    num_heads: int

    def setup(self):
        assert self.d_model % self.num_heads == 0
        self.depth = self.d_model // self.num_heads

        self.wq = nn.Dense(self.d_model)
        self.wk = nn.Dense(self.d_model)
        self.wv = nn.Dense(self.d_model)
        self.dense = nn.Dense(self.d_model)

    def split_heads(self, x, batch_size):
        x = x.reshape((batch_size, -1, self.num_heads, self.depth))
        return x.transpose((0, 2, 1, 3))

    def __call__(self, x):
        batch_size = x.shape[0]
        q = self.split_heads(self.wq(x), batch_size)
        k = self.split_heads(self.wk(x), batch_size)
        v = self.split_heads(self.wv(x), batch_size)

        matmul_qk = jnp.matmul(q, k.transpose((0, 1, 3, 2)))

        dk = jnp.float32(k.shape[-1])
        scaled_attention_logits = matmul_qk / jnp.sqrt(dk)
        attention_weights = nn.softmax(scaled_attention_logits, axis=-1)

        output = jnp.matmul(attention_weights, v)
        output = output.transpose((0, 2, 1, 3)).reshape((batch_size, -1, self.d_model))
        return self.dense(output)

class StandardTransformerBlock(nn.Module):
    d_model: int
    num_heads: int
    dff: int

    def setup(self):
        self.mha = MultiHeadSelfAttention(self.d_model, self.num_heads)
        self.ffn = nn.Sequential([
            nn.Dense(self.dff),
            nn.relu,
            nn.Dense(self.d_model)
        ])

        self.layernorm1 = nn.LayerNorm()
        self.layernorm2 = nn.LayerNorm()

    def __call__(self, x):
        attn_output = self.mha(x)
        out1 = self.layernorm1(x + attn_output)
        ffn_output = self.ffn(out1)
        return self.layernorm2(out1 + ffn_output)

class StandardTransformer(nn.Module):
    d_model: int
    num_heads: int
    num_layers: int
    dff: int
    num_classes: int

    def setup(self):
        self.layers = [StandardTransformerBlock(self.d_model, self.num_heads, self.dff) for _ in range(self.num_layers)]
        self.output_layer = nn.Dense(self.num_classes)

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        x = x.mean(axis=1)  # Reduce the sequence dimension, e.g., by mean pooling
        return self.output_layer(x)

# Define the SaganTransformer again for completeness
class ScaledDotProductAttention(nn.Module):
    d_k: int
    d_model: int  # Add d_model to the constructor

    def setup(self):
        self.query_dense = nn.Dense(self.d_k)
        self.key_dense = nn.Dense(self.d_k)
        self.value_dense = nn.Dense(self.d_k)
        self.output_dense = nn.Dense(self.d_model)  # Change to project back to d_model

    def __call__(self, query, key, value):
        q = self.query_dense(query)
        k = self.key_dense(key)
        v = self.value_dense(value)

        scores = jnp.matmul(q, k.transpose((0, 2, 1))) / jnp.sqrt(self.d_k)
        weights = nn.softmax(scores, axis=-1)
        output = jnp.matmul(weights, v)

        return self.output_dense(output)  # Ensure the output dimension is d_model

class SaganTransformerBlock(nn.Module):
    d_model: int
    n_head: int

    def setup(self):
        self.attention = ScaledDotProductAttention(self.d_model // self.n_head, self.d_model)  # Pass d_model
        self.dense = nn.Dense(self.d_model)
        self.layernorm1 = nn.LayerNorm()
        self.layernorm2 = nn.LayerNorm()
        self.feed_forward = nn.Sequential([nn.Dense(self.d_model * 4), nn.relu, nn.Dense(self.d_model)])

    def __call__(self, x):
        attended = self.attention(x, x, x)
        x = self.layernorm1(x + attended)
        feed_forward_output = self.feed_forward(x)
        x = self.layernorm2(x + feed_forward_output)
        return x

class SaganTransformer(nn.Module):
    d_model: int
    n_head: int
    n_layer: int
    num_classes: int

    def setup(self):
        self.layers = [SaganTransformerBlock(self.d_model, self.n_head) for _ in range(self.n_layer)]
        self.output_layer = nn.Dense(self.num_classes)

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        x = x.mean(axis=1)  # Reduce the sequence dimension, e.g., by mean pooling
        return self.output_layer(x)

# Utility functions for training
def train_step(state, batch):
    def loss_fn(params):
        logits = state.apply_fn({'params': params}, batch['inputs'])
        loss = jnp.mean(optax.softmax_cross_entropy(logits, batch['labels']))
        return loss, logits

    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
    (loss, logits), grads = grad_fn(state.params)
    state = state.apply_gradients(grads=grads)
    accuracy = jnp.mean(jnp.argmax(logits, -1) == jnp.argmax(batch['labels'], -1))
    return state, loss, accuracy

def create_train_state(rng, model, learning_rate):
    params = model.init(rng, jnp.ones([1, 10, model.d_model]))['params']
    tx = optax.adam(learning_rate)
    return train_state.TrainState.create(apply_fn=model.apply, params=params, tx=tx)

def train_and_evaluate(train_data, test_data, model_cls, d_model, num_heads, num_layers, epochs, learning_rate=0.001):
    rng = jax.random.PRNGKey(0)
    model = model_cls(d_model=d_model, n_head=num_heads, n_layer=num_layers)
    state = create_train_state(rng, model, learning_rate)

    loss_history = []
    accuracy_history = []

    for epoch in range(epochs):
        epoch_loss = 0
        epoch_accuracy = 0
        for batch in train_data:
            inputs, labels = batch
            batch = {'inputs': inputs, 'labels': labels}
            state, loss, accuracy = train_step(state, batch)
            epoch_loss += loss
            epoch_accuracy += accuracy
        loss_history.append(epoch_loss / len(train_data))
        accuracy_history.append(epoch_accuracy / len(train_data))

    return loss_history, accuracy_history, state

# Simulated recurrence task dataset
def generate_data(num_samples, sequence_length, d_model, num_classes):
    x = jax.random.normal(random.PRNGKey(0), (num_samples, sequence_length, d_model))
    y = jax.random.randint(random.PRNGKey(1), (num_samples,), 0, num_classes)
    y_one_hot = jax.nn.one_hot(y, num_classes)
    return x, y_one_hot

# Main comparison
def main():
    num_samples = 1000
    sequence_length = 10
    d_model = 64
    num_heads = 8
    num_layers = 2
    epochs = 10
    num_classes = 10

    x, y = generate_data(num_samples, sequence_length, d_model, num_classes)
    train_data = [(x[i:i+32], y[i:i+32]) for i in range(0, num_samples, 32)]
    test_data = train_data  # For simplicity, using the same data for test

    # Train and evaluate both models
    sagan_loss, sagan_acc, sagan_state = train_and_evaluate(train_data, test_data, lambda: SaganTransformer(d_model=d_model, n_head=num_heads, n_layer=num_layers, num_classes=num_classes), d_model, num_heads, num_layers, epochs)
    standard_loss, standard_acc, standard_state = train_and_evaluate(train_data, test_data, lambda: StandardTransformer(d_model=d_model, num_heads=num_heads, num_layers=num_layers, dff=d_model*4, num_classes=num_classes), d_model, num_heads, num_layers, epochs)

    # Plot loss over epochs
    plt.figure(figsize=(12, 6))
    plt.plot(sagan_loss, label='Sagan Transformer')
    plt.plot(standard_loss, label='Standard Transformer')
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.title('Loss over Epochs')
    plt.legend()
    plt.show()

    # Print number of parameters
    def count_params(params):
        return sum(p.size for p in jax.tree_leaves(params))

    print(f"Sagan Transformer Parameters: {count_params(sagan_state.params)}")
    print(f"Standard Transformer Parameters: {count_params(standard_state.params)}")

if __name__ == "__main__":
    main()