import jax
import jax.numpy as jnp
from flax import linen as nn
from flax.training import train_state
import optax
import matplotlib.pyplot as plt
from tqdm import tqdm
import tkinter as tk
from tkinter import filedialog, ttk
import json
from tkinter import messagebox
import traceback
from sklearn.model_selection import train_test_split
import numpy as np

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

    def setup(self):
        self.wq = nn.Dense(self.depth * self.num_heads)
        self.wk = nn.Dense(self.depth * self.num_heads)
        self.wv = nn.Dense(self.depth * self.num_heads)
        self.dense = nn.Dense(self.depth * self.num_heads) 

    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]
        sequence_length = x.shape[1]
        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)

        target_shape = (batch_size, -1, self.depth * self.num_heads)
        output = output.transpose((0, 2, 1, 3)).reshape(target_shape)

        return self.dense(output)

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

    def setup(self):
        self.mha = MultiHeadSelfAttention(depth=self.d_model // self.num_heads, num_heads=self.num_heads, max_output_length=self.max_output_length)
        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):
        try:
            attn_output = self.mha(x)

            # Ensure x and attn_output have the same sequence length
            if x.shape[1] > attn_output.shape[1]:
                x = x[:, -attn_output.shape[1]:]  # Correct slicing for 2D arrays
            elif attn_output.shape[1] > x.shape[1]:
                attn_output = attn_output[:, -x.shape[1]:]  # Correct slicing for 2D arrays

            out1 = self.layernorm1(x + attn_output)
            ffn_output = self.ffn(out1)
            return self.layernorm2(out1 + ffn_output)
        except Exception as e:
            error_message = f"Error occurred during StandardTransformerBlock call: {str(e)}"
            print(error_message)
            traceback.print_exc()
            raise ValueError(error_message) from e
        
class StandardTransformer(nn.Module):
    d_model: int 
    num_heads: int
    num_layers: int
    dff: int
    vocab_size: int
    max_output_length: int

    def setup(self):
        if self.d_model <= 0:
            raise ValueError("d_model must be a positive integer.")
        if self.num_heads <= 0:
            raise ValueError("num_heads must be a positive integer.")
        if self.num_layers <= 0:
            raise ValueError("num_layers must be a positive integer.")
        if self.dff <= 0:
            raise ValueError("dff must be a positive integer.")
        if self.vocab_size <= 0:
            raise ValueError("vocab_size must be a positive integer.")
        if self.max_output_length <= 0:
            raise ValueError("max_output_length must be a positive integer.")

        self.layers = [StandardTransformerBlock(d_model=self.d_model, num_heads=self.num_heads, dff=self.dff, max_output_length=self.max_output_length) for _ in range(self.num_layers)]
        self.output_layer = nn.Dense(self.vocab_size)
        self.embedding = nn.Embed(self.vocab_size, self.d_model)

    def __call__(self, x):
        try:
            if not isinstance(x, (jnp.ndarray, np.ndarray)):
                raise ValueError("Input to StandardTransformer must be a JAX or NumPy array of integers.")
            if not jnp.issubdtype(x.dtype, jnp.integer):
                raise ValueError("Input to StandardTransformer must have an integer dtype.")

            x = self.embedding(x)
            for layer in self.layers:
                x = layer(x)
            return self.output_layer(x)
        except Exception as e:
            error_message = f"Error occurred during StandardTransformer call: {str(e)}"
            print(error_message)
            traceback.print_exc()
            raise ValueError(error_message) from e

# Utility functions for training
def train_step(state, batch):
    try:
        if not isinstance(batch['inputs'], (jnp.ndarray, np.ndarray)):
            raise ValueError("Input batch['inputs'] must be a JAX or NumPy array of integers.")
        if not isinstance(batch['labels'], (jnp.ndarray, np.ndarray)):
            raise ValueError("Input batch['labels'] must be a JAX or NumPy array of integers.")
        if not jnp.issubdtype(batch['inputs'].dtype, jnp.integer):
            raise ValueError("Input batch['inputs'] must have an integer dtype.")
        if not jnp.issubdtype(batch['labels'].dtype, jnp.integer):
            raise ValueError("Input batch['labels'] must have an integer dtype.")

        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
    except Exception as e:
        error_message = f"Error occurred during train_step: {str(e)}"
        print(error_message)
        traceback.print_exc()
        raise ValueError(error_message) from e

def create_train_state(rng, model, learning_rate, vocab_size, max_output_length):
    init_state = model.init(rng, jnp.ones([1, max_output_length], dtype=jnp.int32))
    params = init_state['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, vocab_size, max_output_length, epochs, learning_rate=0.001):
    try:
        if not all(isinstance(batch['inputs'], (jnp.ndarray, np.ndarray)) for batch in train_data):
            raise ValueError("Input 'train_data' must contain batches with JAX or NumPy arrays of integers for 'inputs'.")
        if not all(isinstance(batch['labels'], (jnp.ndarray, np.ndarray)) for batch in train_data):
            raise ValueError("Input 'train_data' must contain batches with JAX or NumPy arrays of integers for 'labels'.")
        if not all(jnp.issubdtype(batch['inputs'].dtype, jnp.integer) for batch in train_data):
            raise ValueError("Input 'train_data' must contain batches with integer dtype for 'inputs'.")
        if not all(jnp.issubdtype(batch['labels'].dtype, jnp.integer) for batch in train_data):
            raise ValueError("Input 'train_data' must contain batches with integer dtype for 'labels'.")
        if not all(isinstance(batch['inputs'], (jnp.ndarray, np.ndarray)) for batch in test_data):
            raise ValueError("Input 'test_data' must contain batches with JAX or NumPy arrays of integers for 'inputs'.")
        if not all(isinstance(batch['labels'], (jnp.ndarray, np.ndarray)) for batch in test_data):
            raise ValueError("Input 'test_data' must contain batches with JAX or NumPy arrays of integers for 'labels'.")
        if not all(jnp.issubdtype(batch['inputs'].dtype, jnp.integer) for batch in test_data):
            raise ValueError("Input 'test_data' must contain batches with integer dtype for 'inputs'.")
        if not all(jnp.issubdtype(batch['labels'].dtype, jnp.integer) for batch in test_data):
            raise ValueError("Input 'test_data' must contain batches with integer dtype for 'labels'.")

        rng = jax.random.PRNGKey(0)
        model = model_cls(d_model=d_model, num_heads=num_heads, num_layers=num_layers, dff=d_model, vocab_size=vocab_size, max_output_length=max_output_length)
        state = create_train_state(rng, model, learning_rate, vocab_size, max_output_length)

        train_loss_history = []
        train_accuracy_history = []

        for epoch in tqdm(range(epochs)):
            epoch_loss = 0
            epoch_accuracy = 0
            for batch in train_data:
                state, loss, accuracy = train_step(state, batch)
                epoch_loss += loss
                epoch_accuracy += accuracy
            train_loss_history.append(epoch_loss / len(train_data))
            train_accuracy_history.append(epoch_accuracy / len(train_data))

            # Evaluate on test data
            test_loss = 0
            test_accuracy = 0
            for batch in test_data:
                inputs, labels = batch
                batch = {'inputs': inputs, 'labels': labels}
                logits = state.apply_fn({'params': state.params}, batch['inputs'])  # Use state.params for evaluation
                loss = jnp.mean(optax.softmax_cross_entropy(logits, batch['labels']))
                test_loss += loss
                accuracy = jnp.mean(jnp.argmax(logits, -1) == jnp.argmax(batch['labels'], -1))
                test_accuracy += accuracy
            test_loss /= len(test_data)
            test_accuracy /= len(test_data)

            print(f"Epoch {epoch+1}: Train Loss: {train_loss_history[-1]}, Train Accuracy: {train_accuracy_history[-1]}, Test Loss: {test_loss}, Test Accuracy: {test_accuracy}")

        return train_loss_history, train_accuracy_history, state, test_loss, test_accuracy 

    except Exception as e:
        error_message = f"Error occurred during train_and_evaluate: {str(e)}"
        print(error_message)
        traceback.print_exc()
        raise ValueError(error_message) from e

def get_input_output_keys(keys):
    possible_input_keys = ['input', 'context', 'source']
    possible_output_keys = ['output', 'response', 'target']

    input_key = next((key for key in keys if key in possible_input_keys), None)
    output_key = next((key for key in keys if key in possible_output_keys), None)

    if input_key is None or output_key is None:
        key_selection_window = tk.Toplevel()
        key_selection_window.title("Select Input and Output Keys")

        input_key_var = tk.StringVar(key_selection_window)
        input_key_menu = ttk.Combobox(key_selection_window, textvariable=input_key_var, values=list(keys))
        input_key_menu.pack()

        output_key_var = tk.StringVar(key_selection_window)
        output_key_menu = ttk.Combobox(key_selection_window, textvariable=output_key_var, values=list(keys))
        output_key_menu.pack()

        def confirm_selection():
            nonlocal input_key, output_key
            input_key = input_key_var.get()
            output_key = output_key_var.get()

            if input_key not in keys or output_key not in keys:
                messagebox.showerror("Error", "Please select valid input and output keys from the dataset.")
                return None, None

            key_selection_window.destroy()
            return input_key, output_key

        confirm_button = tk.Button(key_selection_window, text="Confirm", command=confirm_selection)
        confirm_button.pack()

        key_selection_window.wait_window()
        input_key, output_key = confirm_selection()

    return input_key, output_key

def load_dataset(file_path):
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            lines = [json.loads(line) for line in f]
            if not lines:
                raise ValueError("Dataset is empty.")

            keys = set()
            for line in lines:
                keys.update(line.keys())

        input_key, output_key = get_input_output_keys(keys)

        inputs = []
        outputs = []  # Initialize outputs as a list to store lists of tokens
        vocab = set()

        for line in lines:
            data = line
            inputs.append(data[input_key].split())
            outputs.append(data[output_key].split())  # Append lists of tokens to outputs
            vocab.update(data[input_key].split())
            vocab.update(data[output_key].split())

        # Validate data dimensions before preprocessing
        if not all(isinstance(data, list) for data in inputs):
            raise ValueError("Input 'inputs' must be a list of lists.")
        if not all(isinstance(data, list) for data in outputs):
            raise ValueError("Input 'outputs' must be a list of lists.")
        if not isinstance(vocab, set):
            raise ValueError("Input 'vocab' must be a set.")

        # Call preprocess_data only once inside load_dataset
        inputs, outputs, vocab, max_output_length = preprocess_data(inputs, outputs, vocab)
        inputs = jnp.array(inputs, dtype=jnp.int32)  # Convert inputs to a JAX array
        outputs = jnp.array(outputs)  # Ensure outputs is a 2D array (batch_size, sequence_length)
        token_to_id = {token: i for i, token in enumerate(vocab)}
        id_to_token = {i: token for token, i in token_to_id.items()}
        return inputs, outputs, vocab, max_output_length, token_to_id, id_to_token

    except FileNotFoundError:
        messagebox.showerror("Error", f"File not found: {file_path}")
        print(f"File not found: {file_path}")
        return None, None, None, None, None, None
    except json.JSONDecodeError:
        messagebox.showerror("Error", "Invalid JSONL file format.")
        print("Invalid JSONL file format.")
        return None, None, None, None, None, None
    except KeyError as e:
        messagebox.showerror("Error", f"Key '{e}' not found in the dataset.")
        print(f"Key '{e}' not found in the dataset.")
        return None, None, None, None, None, None
    except Exception as e:
        messagebox.showerror("Error", f"An error occurred while loading the dataset: {str(e)}")
        print(f"An error occurred while loading the dataset: {str(e)}")
        print(traceback.format_exc())
        return None, None, None, None, None, None

def preprocess_data(inputs, outputs, vocab):
    vocab.add('<pad>')  
    token_to_id = {token: i for i, token in enumerate(vocab)}
    id_to_token = {i: token for token, i in token_to_id.items()}

    inputs = [[token_to_id[token] for token in seq] for seq in inputs]
    outputs = [[token_to_id[token] for token in seq] for seq in outputs]

    max_input_length = max(len(seq) for seq in inputs)
    max_output_length = max(len(seq) for seq in outputs)

    inputs = [seq[:max_input_length] + [token_to_id['<pad>']] * (max_input_length - len(seq)) for seq in inputs]
    outputs = [seq[:max_output_length] + [token_to_id['<pad>']] * (max_output_length - len(seq)) for seq in outputs]

    inputs = np.array(inputs, dtype=np.int32)
    outputs = np.array(outputs, dtype=np.int32)

    return inputs, outputs, vocab, max_output_length


def main():
    root = tk.Tk()
    root.withdraw()
    file_path = filedialog.askopenfilename(filetypes=[("JSONL Files", "*.jsonl")])
    if not file_path:
        messagebox.showerror("Error", "No file selected.")
        return

    inputs, outputs, vocab, max_output_length, token_to_id, id_to_token = load_dataset(file_path)
    if inputs is None or outputs is None or vocab is None:
        return

    vocab_size = len(vocab)

    train_inputs, test_inputs, train_outputs, test_outputs = train_test_split(inputs, outputs, test_size=0.2, random_state=42)

    # Convert to JAX arrays
    train_inputs = jnp.array(train_inputs, dtype=jnp.int32)
    test_inputs = jnp.array(test_inputs, dtype=jnp.int32) 

    # Use train_inputs and test_inputs correctly
    train_data = [{'inputs': train_inputs[i], 'labels': train_outputs[i]} for i in range(len(train_inputs))]
    test_data = [{'inputs': test_inputs[i], 'labels': test_outputs[i]} for i in range(len(test_inputs))]

    d_model = 32
    num_heads = 2
    num_layers = 2
    epochs = 2

    train_loss, train_accuracy, state, test_loss, test_accuracy = train_and_evaluate(train_data, test_data, StandardTransformer, d_model, num_heads, num_layers, vocab_size, max_output_length, epochs)

    # Use state, token_to_id, and id_to_token for inference or other purposes
    # ...
    print(f"Final Test Loss: {test_loss}, Final Test Accuracy: {test_accuracy}")

    plt.plot(train_loss, label="Train Loss")
    plt.plot(train_accuracy, label="Train Accuracy")
    plt.xlabel("Epochs")
    plt.ylabel("Loss / Accuracy")
    plt.legend()
    plt.show()

if __name__ == "__main__":
    main()