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
from tkinter import scrolledtext
import re
import traceback

# 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
    vocab_size: int
    max_output_length: 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.vocab_size)
        self.embedding = nn.Embed(self.vocab_size, self.d_model)

    def __call__(self, x):
        x = self.embedding(x)
        for layer in self.layers:
            x = layer(x)
        x = x[:, -self.max_output_length:, :]  # Take only the last max_output_length tokens
        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, vocab_size, max_output_length):
    # Initialize the model
    init_state = model.init(rng, jnp.ones([1, 10, model.d_model])) 

    # Extract the parameters from the init_state dictionary
    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):
    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 = []

    # Use tqdm to track the progress of epochs
    for epoch in tqdm(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
        train_loss_history.append(epoch_loss / len(train_data))
        train_accuracy_history.append(epoch_accuracy / len(train_data))

    return train_loss_history, train_accuracy_history, state

# --- Data Loading and Preprocessing Functions ---
def load_dataset(file_path):
    """Loads and analyzes the dataset from a JSONL file."""
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            # Read the first few lines to analyze keys
            lines = [json.loads(line) for i, line in enumerate(f) if i < 5]
            keys = set()
            for line in lines:
                keys.update(line.keys())

        # Identify input and output keys based on user selection
        input_key, output_key = get_input_output_keys(keys)

        # Create data lists
        inputs = []
        outputs = []
        with open(file_path, 'r', encoding='utf-8') as f:
            for line in f:
                data = json.loads(line)
                inputs.append(data[input_key])
                outputs.append(data[output_key])

        # Preprocess the input and output data
        inputs, outputs, vocab = preprocess_data(inputs, outputs)
        return inputs, outputs, vocab

    except FileNotFoundError:
        messagebox.showerror("Error", f"File not found: {file_path}")
        print(f"File not found: {file_path}")
        return None, None, None
    except json.JSONDecodeError:
        messagebox.showerror("Error", "Invalid JSONL file format.")
        print("Invalid JSONL file format.")
        return 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
    except Exception as e:
        messagebox.showerror("Error", f"An error occurred while loading the dataset: {e}")
        print(f"An error occurred while loading the dataset: {e}")
        traceback.print_exc()  # Print the full traceback to the terminal
        return None, None, None

def preprocess_data(inputs, outputs, max_context_length=1024, max_output_length=512):
    """Preprocesses the input and output data."""
    # Clean the data
    inputs = [clean_text(input) for input in inputs]
    outputs = [clean_text(output) for output in outputs]

    # Add EOS tokens
    outputs = [output + "<eos>" for output in outputs]

    # Create the vocabulary
    vocab = create_vocabulary(inputs, outputs)

    # Chunk the data
    inputs, outputs = chunk_data(inputs, outputs, max_context_length, max_output_length)

    # Convert to numerical data
    inputs = [[vocab[token] for token in input] for input in inputs]
    outputs = [[vocab[token] for token in output] for output in outputs]

    # Pad the data
    inputs = pad_data(inputs, max_context_length, vocab["<pad>"])
    outputs = pad_data(outputs, max_output_length, vocab["<pad>"])

    return inputs, outputs, vocab

def clean_text(text):
    """Cleans the text by removing unsupported characters and punctuation."""
    text = text.lower()
    text = re.sub(r'[^\w\s]', '', text)
    return text

def create_vocabulary(inputs, outputs):
    """Creates a vocabulary from the input and output data."""
    vocab = {"<pad>": 0, "<unk>": 1, "<eos>": 2}
    for input in inputs:
        for token in input.split():
            if token not in vocab:
                vocab[token] = len(vocab)
    for output in outputs:
        for token in output.split():
            if token not in vocab:
                vocab[token] = len(vocab)
    return vocab

def chunk_data(inputs, outputs, max_context_length, max_output_length):
    """Chunks the data into smaller sequences."""
    chunked_inputs = []
    chunked_outputs = []
    for i, input in enumerate(inputs):
        for j in range(0, len(input), max_context_length):
            chunked_inputs.append(input[j:j+max_context_length])
            chunked_outputs.append(outputs[i][j:j+max_output_length])
    return chunked_inputs, chunked_outputs

def pad_data(data, max_length, pad_token):
    """Pads the data to the maximum length."""
    padded_data = []
    for sequence in data:
        padded_sequence = sequence + [pad_token] * (max_length - len(sequence))
        padded_data.append(padded_sequence)
    return padded_data

def get_input_output_keys(keys):
    """Allows the user to choose input and output keys from the dataset."""
    # Create a separate window for key selection
    key_selection_window = tk.Toplevel()  # No need to pass self.master here
    key_selection_window.title("Select Input and Output Keys")

    # Create dropdown menus for input and output keys
    input_key_var = tk.StringVar(key_selection_window)
    input_key_var.set(list(keys)[0])  # Default value
    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_var.set(list(keys)[1])  # Default value
    output_key_menu = ttk.Combobox(key_selection_window, textvariable=output_key_var, values=list(keys))
    output_key_menu.pack()

    # Create a button to confirm the selection
    def confirm_selection():
        input_key = input_key_var.get()
        output_key = output_key_var.get()
        key_selection_window.destroy()
        return input_key, output_key

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

    # Wait for the user to confirm the selection
    key_selection_window.wait_window()
    return confirm_selection()

# --- Model Saving Function ---

def save_model(state):
    """Saves the trained model to a file."""
    file_path = filedialog.asksaveasfilename(defaultextension=".ckpt", filetypes=(("Checkpoint files", "*.ckpt"), ("all files", "*.*")))
    if file_path:
        # Save model state to file 
        jax.experimental.save(file_path, state.params)
        messagebox.showinfo("Model Saved", f"Model saved successfully to: {file_path}")

# --- GUI Implementation ---

class TransformerGUI:
    def __init__(self, master):
        self.master = master
        master.title("Transformer Training GUI")

        # Create frames
        self.dataset_frame = tk.Frame(master)
        self.model_frame = tk.Frame(master)
        self.training_frame = tk.Frame(master)
        self.output_frame = tk.Frame(master)

        # Dataset frame
        self.dataset_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        tk.Label(self.dataset_frame, text="Dataset Settings").pack()

        # Dataset File
        self.dataset_file_label = tk.Label(self.dataset_frame, text="Dataset File:")
        self.dataset_file_label.pack()

        self.dataset_file_entry = tk.Entry(self.dataset_frame, width=50)
        self.dataset_file_entry.pack()

        self.dataset_file_button = tk.Button(self.dataset_frame, text="Browse", command=self.browse_dataset_file)
        self.dataset_file_button.pack()

        # Preview Dataset
        self.preview_button = tk.Button(self.dataset_frame, text="Preview Dataset", command=self.preview_dataset)
        self.preview_button.pack()

        # Model frame
        self.model_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        tk.Label(self.model_frame, text="Model Settings").pack()

        # D Model
        self.d_model_label = tk.Label(self.model_frame, text="D Model:")
        self.d_model_label.pack()

        self.d_model_entry = tk.Entry(self.model_frame, width=10)
        self.d_model_entry.insert(0, "32")
        self.d_model_entry.pack()

        # Num Heads
        self.num_heads_label = tk.Label(self.model_frame, text="Num Heads:")
        self.num_heads_label.pack()

        self.num_heads_entry = tk.Entry(self.model_frame, width=10)
        self.num_heads_entry.insert(0, "4")
        self.num_heads_entry.pack()

        # Num Layers
        self.num_layers_label = tk.Label(self.model_frame, text="Num Layers:")
        self.num_layers_label.pack()

        self.num_layers_entry = tk.Entry(self.model_frame, width=10)
        self.num_layers_entry.insert(0, "8")
        self.num_layers_entry.pack()

        # Max Sequence Length
        self.max_seq_length_label = tk.Label(self.model_frame, text="Max Sequence Length:")
        self.max_seq_length_label.pack()

        self.max_seq_length_entry = tk.Entry(self.model_frame, width=10)
        self.max_seq_length_entry.insert(0, "1024")
        self.max_seq_length_entry.pack()

        # Max Context Length
        self.max_context_length_label = tk.Label(self.model_frame, text="Max Context Length:")
        self.max_context_length_label.pack()

        self.max_context_length_entry = tk.Entry(self.model_frame, width=10)
        self.max_context_length_entry.insert(0, "512")
        self.max_context_length_entry.pack()

        # Max Output Length
        self.max_output_length_label = tk.Label(self.model_frame, text="Max Output Length:")
        self.max_output_length_label.pack()

        self.max_output_length_entry = tk.Entry(self.model_frame, width=10)
        self.max_output_length_entry.insert(0, "256")
        self.max_output_length_entry.pack()

        # Training frame
        self.training_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        tk.Label(self.training_frame, text="Training Settings").pack()

        # Epochs
        self.epochs_label = tk.Label(self.training_frame, text="Epochs:")
        self.epochs_label.pack()

        self.epochs_entry = tk.Entry(self.training_frame, width=10)
        self.epochs_entry.insert(0, "3")
        self.epochs_entry.pack()

        # Learning Rate
        self.learning_rate_label = tk.Label(self.training_frame, text="Learning Rate:")
        self.learning_rate_label.pack()

        self.learning_rate_entry = tk.Entry(self.training_frame, width=10)
        self.learning_rate_entry.insert(0, "0.001")
        self.learning_rate_entry.pack()

        # Start Training Button
        self.start_training_button = tk.Button(self.training_frame, text="Start Training", command=self.start_training)
        self.start_training_button.pack()

        # Save Model Button
        self.save_model_button = tk.Button(self.training_frame, text="Save Model", command=self.save_model)
        self.save_model_button.pack()

        # Output frame
        self.output_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        tk.Label(self.output_frame, text="Output").pack()

        # Text area to display messages
        self.output_text = scrolledtext.ScrolledText(self.output_frame, wrap=tk.WORD, height=5)
        self.output_text.pack()

        # Progress bar
        self.progress_bar = ttk.Progressbar(self.output_frame, orient=tk.HORIZONTAL, length=300, mode="determinate")
        self.progress_bar.pack()

        # Training progress label
        self.progress_label = tk.Label(self.output_frame, text="")
        self.progress_label.pack()

    def browse_dataset_file(self):
        """Opens a file dialog to select the dataset file."""
        file_path = filedialog.askopenfilename(initialdir="/", title="Select Dataset File",
                                              filetypes=(("JSONL files", "*.jsonl"), ("all files", "*.*")))
        self.dataset_file_entry.delete(0, tk.END)
        self.dataset_file_entry.insert(0, file_path)

    def preview_dataset(self):
        """Opens a new window to preview the selected dataset file."""
        file_path = self.dataset_file_entry.get()
        if not file_path:
            self.show_error("Error", "Please select a dataset file.")
            return

        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                # Read the first few lines to analyze keys
                lines = [json.loads(line) for i, line in enumerate(f) if i < 5]
                keys = set()
                for line in lines:
                    keys.update(line.keys())

            # Create a new window for the preview
            preview_window = tk.Toplevel(self.master)
            preview_window.title("Dataset Preview")

            # Create a table to display the data
            preview_table = ttk.Treeview(preview_window, columns=list(keys), show="headings")
            for key in keys:
                preview_table.heading(key, text=key)
            preview_table.pack()

            # Add the data to the table
            for i, line in enumerate(lines):
                values = [line[key] for key in keys]
                preview_table.insert("", "end", values=values)

        except FileNotFoundError:
            self.show_error("Error", "File not found: " + file_path)
        except json.JSONDecodeError:
            self.show_error("Error", "Invalid JSONL file format.")
        except Exception as e:
            self.show_error("Error", f"An error occurred while previewing the dataset: {e}")

    def start_training(self):
        """Starts the training process."""
        try:
            # Get model parameters
            d_model = int(self.d_model_entry.get())
            num_heads = int(self.num_heads_entry.get())
            num_layers = int(self.num_layers_entry.get())
            max_seq_length = int(self.max_seq_length_entry.get())
            max_context_length = int(self.max_context_length_entry.get())
            max_output_length = int(self.max_output_length_entry.get())
            epochs = int(self.epochs_entry.get())
            learning_rate = float(self.learning_rate_entry.get())

            # Load dataset
            file_path = self.dataset_file_entry.get()
            if not file_path:
                self.show_error("Error", "Please select a dataset file.")
                return

            train_data, test_data, vocab = load_dataset(file_path)
            if train_data is None or test_data is None:
                return  # Dataset loading failed

            # Start training
            self.show_message("Info", "Training Transformer...")
            train_loss_history, train_accuracy_history, state = train_and_evaluate(
                train_data, test_data, StandardTransformer, d_model, num_heads, num_layers, len(vocab), max_output_length,
                epochs, learning_rate
            )

            # Save the trained model
            self.save_model(state)  # Make sure to pass 'state' to save_model

            # Plot training progress
            plt.figure(figsize=(12, 6))
            plt.plot(train_loss_history, label="Training Loss")
            plt.xlabel('Epochs')
            plt.ylabel('Loss')
            plt.title('Training Loss over Epochs')
            plt.legend()
            plt.show()

            self.show_message("Success", "Training finished successfully.")

        except ValueError as e:
            self.show_error("Error", f"Invalid input for parameters: {e}")
            print(f"Invalid input for parameters: {e}")
            traceback.print_exc()  # Print the full traceback to the terminal
        except Exception as e:  # Catch all exceptions
            self.show_error("Error", f"An error occurred during training: {e}")
            print(f"An error occurred during training: {e}")
            traceback.print_exc()  # Print the full traceback to the terminal

    def save_model(self, state):
        """Saves the trained model to a file."""
        file_path = filedialog.asksaveasfilename(defaultextension=".ckpt", filetypes=(("Checkpoint files", "*.ckpt"), ("all files", "*.*")))
        if file_path:
            # Save model state to file 
            jax.experimental.save(file_path, state.params)
            self.show_message("Success", f"Model saved successfully to: {file_path}")

    def show_error(self, title, message):
        """Displays an error message box and logs the error to the terminal."""
        messagebox.showerror(title, message)
        self.show_message("Error", message)
        print(f"{title}: {message}")

    def show_message(self, title, message):
        """Displays a message box and appends the message to the output text area."""
        messagebox.showinfo(title, message)
        self.output_text.config(state=tk.NORMAL)  # Enable text editing
        self.output_text.insert(tk.END, f"{title}: {message}\n")
        self.output_text.config(state=tk.DISABLED)  # Disable text editing

# Create the main window and start the GUI
root = tk.Tk()
app = TransformerGUI(root)
root.mainloop()