import os
import numpy as np
from Bio import SeqIO
from sklearn.model_selection import train_test_split
from tensorflow.keras.preprocessing.text import Tokenizer
from tensorflow.keras.preprocessing.sequence import pad_sequences

def preprocess_dataset(dataset_path, max_sequence_length):
    sequences = []
    labels = []
    classes = os.listdir(dataset_path)
    
    for class_name in classes:
        class_path = os.path.join(dataset_path, class_name)
        for file_name in os.listdir(class_path):
            file_path = os.path.join(class_path, file_name)
            sequence = str(SeqIO.read(file_path, "fasta").seq)
            sequences.append(sequence)
            labels.append(class_name)
    
    tokenizer = Tokenizer()
    tokenizer.fit_on_texts(sequences)
    sequences = tokenizer.texts_to_sequences(sequences)
    sequences = pad_sequences(sequences, maxlen=max_sequence_length)
    
    label_mapping = {class_name: i for i, class_name in enumerate(set(labels))}
    labels = [label_mapping[label] for label in labels]
    
    sequences = np.array(sequences)
    labels = np.array(labels)
    
    train_sequences, test_sequences, train_labels, test_labels = train_test_split(
        sequences, labels, test_size=0.2, random_state=42)
    
    return train_sequences, test_sequences, train_labels, test_labels, label_mapping