# %%

__author__ = "Sarah Shi"

import os
import math
import time
import copy
import pickle
import warnings

import numpy as np
import pandas as pd

from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, confusion_matrix

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torch.nn.functional as F

from .core import *
from .core import same_seeds
# same_seeds(2)
from .stoichiometry import *
from .constants import OXIDES

import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import matplotlib.cm as mcm
from sklearn.decomposition import PCA


# %%

from sklearn.base import BaseEstimator, TransformerMixin

class CLRTransformer(BaseEstimator, TransformerMixin):
    """
    Proper CLR:
      - optional row renormalization to a constant sum
      - eps-clipping to avoid log(0)
      - per-sample (row-wise) log-centering
      - preserves column names and can return a DataFrame
    """
    def __init__(self, eps=1e-5, renorm_to=100.0, clip=True, return_dataframe=True):
        self.eps = eps
        self.renorm_to = renorm_to
        self.clip = clip
        self.return_dataframe = return_dataframe
        self.feature_names_in_ = None
        self.output_feature_names_ = None

    def fit(self, X, y=None):
        if isinstance(X, pd.DataFrame):
            self.feature_names_in_ = X.columns.tolist()
        else:
            self.feature_names_in_ = None
        # CLR has no learned parameters
        if self.feature_names_in_ is not None:
            self.output_feature_names_ = [f"CLR({c}/G)" for c in self.feature_names_in_]
        return self

    def transform(self, X):
        is_df = isinstance(X, pd.DataFrame)
        cols = self.feature_names_in_ if (is_df and self.feature_names_in_ is not None) else (
            X.columns.tolist() if is_df else None
        )
        A = X.values if is_df else np.asarray(X, dtype=float)

        # optional renormalization to constant sum
        if self.renorm_to is not None:
            row_sums = A.sum(axis=1, keepdims=True)
            row_sums[row_sums == 0] = np.nan
            A = (A / row_sums) * self.renorm_to

        # strictly positive
        if self.clip:
            A = np.clip(A, self.eps, None)
        else:
            if np.any(A <= 0):
                raise ValueError("CLR requires strictly positive inputs. Set clip=True or pre-clean zeros.")

        logA = np.log(A)
        clrA = logA - logA.mean(axis=1, keepdims=True)  # <-- row-wise geometric-mean centering

        if self.return_dataframe and cols is not None:
            out_cols = self.output_feature_names_ or [f"CLR({c}/G)" for c in cols]
            return pd.DataFrame(clrA, index=(X.index if is_df else None), columns=out_cols)
        return clrA

    def save(self, filepath):
        with open(filepath, 'wb') as f:
            pickle.dump({
                'eps': self.eps,
                'renorm_to': self.renorm_to,
                'clip': self.clip,
                'return_dataframe': self.return_dataframe,
                'feature_names_in_': self.feature_names_in_,
            }, f)

    def load(cls, filepath):
        with open(filepath, 'rb') as f:
            d = pickle.load(f)
        obj = cls(eps=d['eps'], renorm_to=d['renorm_to'], clip=d['clip'],
                  return_dataframe=d['return_dataframe'])
        obj.feature_names_in_ = d['feature_names_in_']
        if obj.feature_names_in_ is not None:
            obj.output_feature_names_ = [f"CLR({c}/G)" for c in obj.feature_names_in_]
        return obj


class MultiClassClassifier(nn.Module):

    """
    A neural network module for multi-class classification tasks. It
    consists of a sequence of layers defined by the input dimensions, number
    of classes, dropout rate, and sizes of hidden layers. It can be
    customized with different numbers and sizes of hidden layers, as well as
    varying dropout rates to prevent overfitting. The final output layer is
    designed for classification among a fixed number of classes.

    Parameters:
        input_dim (int): Dimensionality of the input features. Defaults to 11.
        classes (int): The number of output classes for classification. Defaults to 24.
        dropout_rate (float): The dropout rate applied after each hidden layer. Defaults to 0.1.
        hidden_layer_sizes (list of int): The sizes of each hidden layer. Defaults to a single
                                          hidden layer with 8 units.

    Attributes:
        input_dim (int): Internal storage of the input dimensionality.
        classes (int): Internal storage of the number of classes.
        dropout_rate (float): Internal storage of the dropout rate.
        hls (list of int): Internal storage of the hidden layer sizes.
        encode (nn.Sequential): The sequential container of layers making up the encoder part
                                of the classifier, including linear, batch normalization,
                                leaky ReLU, and dropout layers.

    Methods:
        encoded(x): Encodes input `x` through the sequence of layers defined in `encode`.
        forward(x): Implements the forward pass of the network, returning raw scores for each class.
        predict(x): Provides class predictions for input `x` based on the scores from the forward pass.

    The class utilizes a helper function `element` to create each hidden layer or the variational
    layer if it is the last one. The `weights_init` function is applied to initialize weights
    after the model is constructed.

    """

    def __init__(
        self,
        input_dim=11,
        classes=23,
        dropout_rate=0.1,
        hidden_layer_sizes=[64, 32, 16],
    ):
        super(MultiClassClassifier, self).__init__()
        self.input_dim = input_dim
        self.classes = classes
        self.dropout_rate = dropout_rate
        self.hls = hidden_layer_sizes

        def element(in_channel, out_channel, is_last=False):
            if not is_last:
                layers = [
                    nn.Linear(in_channel, out_channel),
                    nn.BatchNorm1d(out_channel), # Add batch normalization
                    nn.LeakyReLU(0.02),
                    nn.Dropout(self.dropout_rate), # Add dropout
                ]
            else:
                layers = [VariationalLayer(in_channel, out_channel)]
            return layers

        encoder = []
        for i, size in enumerate(self.hls):
            if i == 0:
                encoder += element(
                    self.input_dim, size, is_last=(i == len(self.hls) - 1)
                )
            else:
                encoder += element(
                    self.hls[i - 1], size, is_last=(i == len(self.hls) - 1)
                )

        encoder += [nn.Linear(size, self.classes)]  # Add this line

        self.encode = nn.Sequential(*encoder)
        self.apply(weights_init)

    def encoded(self, x):
        return self.encode(x)

    def forward(self, x):
        en = self.encoded(x)
        return en

    def predict(self, x):
        # Get predicted scores
        scores = self.forward(x)
        # Get predicted class indices
        class_indices = scores.argmax(dim=1)
        return class_indices


def predict_class_prob_nn_train(model, input_data, n_iterations=50):
    """

    Computes the predicted class probabilities for the given input data using the model by
    performing multiple forward passes. The function operates in evaluation mode and does not
    track gradients. It returns the mean and standard deviation of the softmax probabilities
    across all iterations, providing a measure of model uncertainty.

    Parameters:
        model (nn.Module): The model to be used for prediction, which should already be trained.
        input_data (Tensor): The input data to be passed to the model for prediction.
        n_iterations (int): The number of forward passes to perform for prediction. Defaults to 100.

    Returns:
        prediction_mean (ndarray): The mean class probabilities across all iterations.
        prediction_std (ndarray): The standard deviation of class probabilities, indicating uncertainty.

    """

    model.eval()
    output_list = []
    for i in range(n_iterations):
        with torch.no_grad():
            output = model(input_data)
            output_list.append(
                torch.nn.functional.softmax(output, dim=1).detach().cpu().numpy()
            )

    output_list = np.array(output_list)

    # Calculate mean and standard deviation
    prediction_mean = output_list.mean(axis=0)
    prediction_std = output_list.std(axis=0)

    return prediction_mean, prediction_std


def predict_class_prob_nn(df, n_iterations=50):
    """

    Predicts the class probabilities, corresponding mineral names, and the maximum
    probability for each class using a predefined MultiClassClassifier model. This
    function loads a pre-trained model and its optimizer state, normalizes input
    data, and performs multiple inference iterations to compute the prediction probabilities.

    Parameters:
        df (DataFrame): The input DataFrame containing the oxide composition data.
        n_iterations (int): The number of inference iterations to average over for predictions.

    Returns:
        df (DataFrame): The input DataFrame with columns predict_mineral (predicted mineral names)
        and predict_prob (maximum probability of predicted class).
        probability_matrix (ndarray): The matrix of class probabilities for each sample.

    """
    
    # Create output DataFrame with original indices and structure
    result_df = df.copy()
    pred_cols = ["Predict_Mineral", "Predict_Probability", 
                 "Second_Predict_Mineral", "Second_Predict_Probability"]
    
    for col in pred_cols:
        result_df[col] = np.nan
    
    # Initialize with proper dtypes
    result_df["Predict_Mineral"] = pd.Series(index=df.index, dtype="object")
    result_df["Second_Predict_Mineral"] = pd.Series(index=df.index, dtype="object")
    result_df["Predict_Probability"] = pd.Series(index=df.index, dtype="float64")
    result_df["Second_Predict_Probability"] = pd.Series(index=df.index, dtype="float64")

    # --- detect rows where *every* input oxide value is 0 ---
    oxide_cols_in_df = [c for c in OXIDES if c in df.columns]
    if oxide_cols_in_df:
        # treat NaN as 0 for this check so [0, NaN, 0] counts as "all zero"
        all_zero_mask = (df[oxide_cols_in_df].fillna(0).to_numpy() == 0).all(axis=1)
        all_zero_mask = pd.Series(all_zero_mask, index=df.index)
    else:
        all_zero_mask = pd.Series(False, index=df.index)

    # explicitly set outputs (minerals already NaN, but make it explicit)
    result_df.loc[all_zero_mask, "Predict_Mineral"] = np.nan
    result_df.loc[all_zero_mask, "Second_Predict_Mineral"] = np.nan
    result_df.loc[all_zero_mask, "Predict_Probability"] = np.nan
    result_df.loc[all_zero_mask, "Second_Predict_Probability"] = np.nan
    # ------------------------------------------------------------

    # Identify and classify zircons
    zircon_mask = (df['ZrO2'] > 50) if 'ZrO2' in df.columns else pd.Series(False, index=df.index)
    zircon_mask = zircon_mask & ~all_zero_mask
    result_df.loc[zircon_mask, "Predict_Mineral"] = "Zircon"
    result_df.loc[zircon_mask, "Predict_Probability"] = 1.0
    result_df.loc[zircon_mask, "Second_Predict_Mineral"] = np.nan
    result_df.loc[zircon_mask, "Second_Predict_Probability"] = np.nan

    # Process non-zircon samples
    non_zircon_mask = (~zircon_mask) & (~all_zero_mask)
    probability_matrix = np.array([])

    if non_zircon_mask.any():
        non_za_df = df[non_zircon_mask]
        
        # Neural network setup — model and device
        oxides = OXIDES
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        model = MultiClassClassifier(
            input_dim=len(oxides),
            dropout_rate=0.1,
            hidden_layer_sizes=[64, 32, 16]
        ).to(device)
        
        # set optimizer to None for inference
        optimizer = None # torch.optim.SGD(model.parameters(), lr=5e-3, weight_decay=1e-3)
        model_path = os.path.join(os.path.dirname(__file__), "nn_best_model_v0019.pt")
        load_model(model, optimizer, model_path)
        
        # Normalize data
        norm_wt = norm_data_nn(non_za_df).astype(np.float32, copy=False)
        input_data = torch.from_numpy(norm_wt).to(device) # torch.from_numpy better than torch.Tensor

        model.eval()
        with torch.inference_mode():
            device = next(model.parameters()).device
            N = len(input_data)
            C = model.classes
            BATCH = 2**13 
            K = 10 # MC passes per batch chunk (grouped to cut Python overhead)

            probs_mean = torch.empty((N, C), device=device, dtype=torch.float32)

            for start in range(0, N, BATCH):
                end = min(start + BATCH, N)
                x = input_data[start:end] 
                b = x.shape[0]

                done = 0
                acc = torch.zeros((b, C), device=device, dtype=torch.float32)

                while done < n_iterations:
                    kk = min(K, n_iterations - done)

                    # run kk independent forwards and average on the device
                    outs = []
                    for _ in range(kk):
                        logits = model(x)
                        outs.append(torch.softmax(logits, dim=1))
                    acc += torch.stack(outs, dim=0).mean(dim=0) * kk
                    done += kk

                probs_mean[start:end] = acc / float(n_iterations)

            probability_matrix = probs_mean.detach().cpu().numpy()

        # Get top predictions efficiently
        top_two_indices = np.argsort(probability_matrix, axis=1)[:, -2:]
        first_probs = probability_matrix[np.arange(len(probability_matrix)), top_two_indices[:, 1]]
        second_probs = probability_matrix[np.arange(len(probability_matrix)), top_two_indices[:, 0]]
        first_mins = class2mineral_nn(top_two_indices[:, 1])
        second_mins = class2mineral_nn(top_two_indices[:, 0])

        # Update result dataframe
        result_df.loc[non_zircon_mask, "Predict_Mineral"] = first_mins
        result_df.loc[non_zircon_mask, "Predict_Probability"] = first_probs
        result_df.loc[non_zircon_mask, "Second_Predict_Mineral"] = second_mins
        result_df.loc[non_zircon_mask, "Second_Predict_Probability"] = second_probs

    # Process specialized classifiers (unchanged, but could also be optimized)
    oxide_cols = [c for c in result_df.columns if c in OXIDES]
    mineral_col = "Predict_Mineral" if "Predict_Mineral" in result_df.columns else None
    cols = oxide_cols + ([mineral_col] if mineral_col else [])

    def _merge_subclass(mask, Classifier, want_sub=True):
        if not mask.any():
            return
        sub = result_df.loc[mask, cols]
        clf = Classifier(sub)
        # Ensure subclass=True to get submineral back by default
        out = clf.classify(subclass=want_sub)
        # Expect out to have "Mineral" and (if want_sub) "Submineral"
        if "Mineral" in out.columns:
            result_df.loc[mask, "Predict_Mineral"] = out["Mineral"].values
        if want_sub and "Submineral" in out.columns:
            result_df.loc[mask, "Submineral"] = out["Submineral"].values

    # Pyroxene classification
    px_mask = result_df["Predict_Mineral"] == "Pyroxene"
    _merge_subclass(px_mask, PyroxeneClassifier, want_sub=True)

    # Feldspar classification
    fspar_mask = result_df["Predict_Mineral"] == "Feldspar"
    _merge_subclass(fspar_mask, FeldsparClassifier, want_sub=True)
    
    # Oxide classification
    ox_mask = result_df["Predict_Mineral"].isin(["Rhombohedral_Oxides", "Spinels"])
    _merge_subclass(ox_mask, OxideClassifier, want_sub=True)

    if "Submineral" not in result_df.columns:
        result_df["Submineral"] = pd.Series(index=result_df.index, dtype="object")

    cols = list(result_df.columns)
    if "Predict_Mineral" in cols and "Submineral" in cols:
        # remove and re-insert Submineral right after Predict_Mineral
        cols.remove("Submineral")
        insert_at = cols.index("Predict_Mineral") + 1
        cols.insert(insert_at, "Submineral")
        result_df = result_df[cols]

    return result_df, probability_matrix


def train_nn(
    model,
    optimizer,
    train_loader,
    valid_loader,
    n_epoch,
    criterion,
    kl_weight_decay,
    kl_decay_epochs=750,
    patience=50,
):
    """

    Trains a neural network model using the provided data loaders, optimizer, and loss criterion. It incorporates KL divergence
    into the loss to enable learning in a variational framework, with the KL weight increasing each epoch until a maximum value
    is reached. The function includes an early stopping mechanism that terminates training if validation loss does not improve
    for a specified number of consecutive epochs.

    Parameters:
        model (nn.Module): The neural network model to train.
        optimizer (Optimizer): The optimization algorithm used to update model weights.
        train_loader (DataLoader): The DataLoader containing the training data.
        valid_loader (DataLoader): The DataLoader containing the validation data.
        n_epoch (int): The total number of training epochs.
        criterion (Loss): The loss function used for training.
        kl_weight_decay (float): The increment to the KL divergence weight per epoch.
        kl_decay_epochs (int): The number of epochs over which to increment the KL weight. Defaults to 750.
        patience (int): The number of epochs to wait for improvement in validation loss before early stopping. Defaults to 50.

    Returns:
        train_output (Tensor): The output from the model for the last training batch.
        valid_output (Tensor): The output from the model for the last validation batch.
        avg_train_loss (list): The list of average training losses per epoch.
        avg_valid_loss (list): The list of average validation losses per epoch.
        best_valid_loss (float): The best validation loss observed during training.
        best_model_state (dict): The state dictionary of the model at the point of the best validation loss.

    """

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    avg_train_loss = []
    avg_valid_loss = []
    best_valid_loss = float("inf")
    best_model_state = None
    patience_counter = 0

    kl_weight = 0.0  # Initial kl_weight

    for epoch in range(n_epoch):
        model.train()
        t = time.time()
        train_loss = []
        for i, (data, labels) in enumerate(train_loader):
            x = data.to(device)
            y = labels.to(device)
            train_output = model(x)
            loss = criterion(train_output, y)

            # Add KL divergence with weight decay
            kl_div = 0.0
            # kl_weight = min(kl_weight + (kl_weight_decay * (epoch // kl_decay_epochs)), 1)
            kl_weight_increment = kl_weight_decay / kl_decay_epochs
            kl_weight = min(kl_weight + kl_weight_increment, 1)

            for module in model.modules():
                if isinstance(module, VariationalLayer):
                    kl_div += module.kl_divergence()
            loss += kl_weight * kl_div / len(train_loader.dataset)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            train_loss.append(loss.detach().item())

        # Validation
        model.eval()
        valid_loss = []
        with torch.no_grad():
            for i, (data, labels) in enumerate(valid_loader):
                x = data.to(device)
                y = labels.to(device)
                valid_output = model(x)
                loss = criterion(valid_output, y)
                valid_loss.append(loss.detach().item())

        # Logging
        avg_train = sum(train_loss) / len(train_loss)
        avg_valid = sum(valid_loss) / len(valid_loss)
        avg_train_loss.append(avg_train)
        avg_valid_loss.append(avg_valid)

        training_time = time.time() - t

        print(
            f"[{epoch+1:03}/{n_epoch:03}] train_loss: {avg_train:.6f}, valid_loss: {avg_valid:.6f}, time: {training_time:.2f} s"
        )

        # Early stopping
        if avg_valid < best_valid_loss:
            best_valid_loss = avg_valid
            patience_counter = 0
            best_model_state = copy.deepcopy(
                model.state_dict()
            )  # Save the best model weights

        else:
            patience_counter += 1
            if patience_counter >= patience:
                print(
                    f"Validation loss hasn't improved for {patience} epochs. Stopping early."
                )
                break

    return (
        train_output,
        valid_output,
        avg_train_loss,
        avg_valid_loss,
        best_valid_loss,
        best_model_state,
    )


def neuralnetwork(df, hls_list, kl_weight_decay_list, lr, wd, dr, ep, n, balanced):
    """

    Trains a neural network with various configurations of hidden layer sizes and KL weight
    decay parameters to find the best model for classifying minerals based on their oxide
    composition. It normalizes input data, balances the dataset if required, initializes
    the model and optimizer, and performs training and validation. The best performing
    model's parameters are saved, along with training and validation losses, and prediction
    reports.

    Parameters:
        df (DataFrame): The input DataFrame with mineral composition data and labels.
        hls_list (list of list of int): List of configurations for hidden layer sizes.
        kl_weight_decay_list (list of float): List of KL weight decay values to try during training.
        lr (float): Learning rate for the optimizer.
        wd (float): Weight decay factor for regularization.
        dr (float): Dropout rate for the model.
        ep (int): Number of epochs to train.
        n (float): Test size fraction or absolute number for splitting the dataset.
        balanced (bool): Whether to balance the dataset or not.

    Returns:
        best_model_state (dict): The state dictionary of the best performing model.

    """

    path_beg = os.getcwd() + "/"
    output_dir = ["parametermatrix_neuralnetwork"]
    for ii in range(len(output_dir)):
        if not os.path.exists(path_beg + output_dir[ii]):
            os.makedirs(path_beg + output_dir[ii], exist_ok=True)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Split the dataset into train and test sets
    train_df, valid_df = train_test_split(df, test_size=n, random_state=42)
    if balanced == True:
        train_df = balance(train_df, n=1000)
    train_df.to_csv('train_df_nn.csv', index=False)

    for _df in (train_df, valid_df):
        _df['Mineral'] = (
            _df['Mineral'].astype(str)
            .replace(['Clinopyroxene', 'Orthopyroxene'], 'Pyroxene')
            .replace(['Plagioclase', 'KFeldspar'], 'Feldspar')
            .replace(['Hematite', 'Ilmenite'], 'Rhombohedral_Oxides')
            .replace(['Magnetite', 'Spinel'], 'Spinels')
    )
    
    train_df_nozirc = train_df[train_df['Mineral'] != 'Zircon'].copy()
    valid_df_nozirc = valid_df[valid_df['Mineral'] != 'Zircon'].copy()
            
    all_cats = pd.Categorical(train_df_nozirc["Mineral"])
    mapping = dict(enumerate(all_cats.categories))

    inv_mapping = {cat: idx for idx, cat in mapping.items()}
    sort_mapping = dict(sorted(mapping.items(), key=lambda item: item[0]))

    train_df_nozirc["_code"] = train_df_nozirc["Mineral"].map(inv_mapping).astype(int)
    valid_df_nozirc["_code"] = valid_df_nozirc["Mineral"].map(inv_mapping).astype(int)


    # ========== CORRECTED CLR IMPLEMENTATION ==========
    # Use CLRTransformer
    # clr = CLRTransformer(eps=1e-6, renorm_to=None, clip=True, return_dataframe=False)
    
    # Fill NaNs with 0 first (CLR will clip them to eps)
    # train_filled = train_df_nozirc[OXIDES].fillna(0)
    # valid_filled = valid_df_nozirc[OXIDES].fillna(0)
    
    # Apply CLR transformation
    # train_x_clr = clr.fit_transform(train_filled)
    # valid_x_clr = clr.transform(valid_filled)
    
    # Optional: Scale the CLR-transformed data
    # ss = StandardScaler().fit(train_x_clr)
    # train_x = ss.transform(train_x_clr)
    # valid_x = ss.transform(valid_x_clr)

    # Save both transformers
    # clr_path = os.path.join(path_beg, "parametermatrix_neuralnetwork", "clr_transformer_v001.pkl")
    # clr.save(clr_path)
    
    # scaler_path = os.path.join(path_beg, "parametermatrix_neuralnetwork", "clr_scaler_nn_v001.npz")
    # np.savez(scaler_path,
    #          mean=ss.mean_,
    #          scale=ss.scale_,
    #          feature_names=OXIDES)

    # scale
    ss = StandardScaler().fit(train_df_nozirc[OXIDES].fillna(0))
    train_x = ss.transform(train_df_nozirc[OXIDES].fillna(0))
    valid_x = ss.transform(valid_df_nozirc[OXIDES].fillna(0))
    scaler_path = os.path.join(path_beg, "parametermatrix_neuralnetwork", "scaler_nn_v0019.npz")
    np.savez(scaler_path,
             mean=pd.Series(ss.mean_, index=OXIDES),
             scale=pd.Series(np.sqrt(ss.var_), index=OXIDES))

    # encode labels
    train_y = train_df_nozirc["_code"].to_numpy()
    valid_y = valid_df_nozirc["_code"].to_numpy()

    # Define datasets to be used with PyTorch - see autoencoder file for details
    feature_dataset = LabelDataset(train_x, train_y)
    valid_dataset = LabelDataset(valid_x, valid_y)

    # Autoencoder params:
    lr = lr
    wd = wd
    dr = dr
    epochs = ep
    batch_size = 256
    input_size = len(feature_dataset.__getitem__(0)[0])

    best_hidden_layer_size = None
    best_kl_weight_decay = None
    best_model_state = None
    best_valid_loss = float("inf")

    # Define data loaders
    feature_loader = DataLoader(feature_dataset, batch_size=batch_size, shuffle=True)
    valid_loader = DataLoader(valid_dataset, batch_size=batch_size) #, shuffle=True)
    np.savez(
        "parametermatrix_neuralnetwork/" + str(lr) + "_" + str(wd) + "_" + str(dr) + "_" + str(ep) + "_best_model_nn_features_clr.npz",
        feature_loader=feature_loader,
        valid_loader=valid_loader,
    )

    train_losses_dict = {}
    valid_losses_dict = {}

    for hls in hls_list:
        for kl_weight_decay in kl_weight_decay_list:
            print(
                f"Training with KL weight decay: {kl_weight_decay} and hidden layer sizes: {hls}"
            )

            # Initialize model
            model = MultiClassClassifier(
                input_dim=input_size, dropout_rate=dr, hidden_layer_sizes=hls
            ).to(device)

            # Define loss function and optimizer
            criterion = nn.CrossEntropyLoss()
            optimizer = torch.optim.SGD(model.parameters(), lr=lr, weight_decay=wd)

            # Train model and get the best test loss and model state
            (
                train_output,
                valid_output,
                avg_train_loss,
                avg_valid_loss,
                current_best_valid_loss,
                current_best_model_state,
            ) = train_nn(
                model,
                optimizer,
                feature_loader,
                valid_loader,
                epochs,
                criterion,
                kl_weight_decay=kl_weight_decay,
            )

            if current_best_valid_loss < best_valid_loss:
                best_valid_loss = current_best_valid_loss
                best_kl_weight_decay = kl_weight_decay
                best_hidden_layer_size = hls
                best_model_state = current_best_model_state

            train_losses_dict[(kl_weight_decay, tuple(hls))] = avg_train_loss
            valid_losses_dict[(kl_weight_decay, tuple(hls))] = avg_valid_loss

    # Create a new model with the best model state
    best_model = MultiClassClassifier(
        input_dim=input_size, dropout_rate=dr, hidden_layer_sizes=best_hidden_layer_size
    )
    best_model.load_state_dict(best_model_state)
    best_model.eval()

    # Perform predictions on the test dataset using the best model
    with torch.no_grad():
        valid_predictions = best_model(torch.Tensor(valid_x))
        train_predictions = best_model(torch.Tensor(train_x))
        valid_pred_y = valid_predictions.argmax(dim=1).cpu().numpy()
        train_pred_y = train_predictions.argmax(dim=1).cpu().numpy()

    # Calculate classification metrics for the test dataset
    valid_report = classification_report(
        valid_y,
        valid_pred_y,
        target_names=list(sort_mapping.values()),
        zero_division=0,
        output_dict=True,
    )
    train_report = classification_report(
        train_y,
        train_pred_y,
        target_names=list(sort_mapping.values()),
        zero_division=0,
        output_dict=True,
    )

    # Print the best kl_weight_decay value and test report
    print("Best kl_weight_decay:", best_kl_weight_decay)
    print("Best best_hidden_layer_size:", best_hidden_layer_size)

    # Save the best model and other relevant information
    model_path = "parametermatrix_neuralnetwork/" + str(lr) + "_" + str(wd) + "_" + str(dr) + "_" + str(ep) + "_" + str(kl_weight_decay) + "_" + str(hls) + "_best_model.pt"
    save_model_nn(optimizer, best_model_state, model_path)

    train_pred_mean, train_pred_std = predict_class_prob_nn_train(
        model, feature_dataset.x
    )
    valid_pred_mean, valid_pred_std = predict_class_prob_nn_train(
        model, valid_dataset.x
    )

    # Get the most probable classes
    train_pred_y = np.argmax(train_pred_mean, axis=1)
    valid_pred_y = np.argmax(valid_pred_mean, axis=1)

    np.savez(
        "parametermatrix_neuralnetwork/" + str(lr) + "_" + str(wd) + "_" + str(dr) + "_" + str(ep) + "_" + str(kl_weight_decay) + "_" + str(hls) + "_best_model_data.npz",
        best_hidden_layer_size=best_hidden_layer_size,
        best_kl_weight_decay=best_kl_weight_decay,
        valid_report=valid_report,
        train_report=train_report,
        train_y=train_y,
        valid_y=valid_y,
        train_pred_y=train_pred_y,
        valid_pred_y=valid_pred_y,
        train_pred_mean=train_pred_mean,
        train_pred_std=train_pred_std,
        valid_pred_mean=valid_pred_mean,
        valid_pred_std=valid_pred_std,
    )

    # Save the train and test losses
    np.savez(
        "parametermatrix_neuralnetwork/" + str(lr) + "_" + str(wd) + "_" + str(dr) + "_" + str(ep) + "_best_model_losses.npz",
        train_losses=train_losses_dict,
        valid_losses=valid_losses_dict,
    )

    return best_model_state