"""
Argument Parser for Drug Repositioning Module

This module defines command-line arguments for configuring the drug-disease association prediction
training process. The arguments include general settings, training parameters, and model hyperparameters.
"""

import argparse

parser = argparse.ArgumentParser(
    formatter_class=argparse.ArgumentDefaultsHelpFormatter
)

# General arguments
parser.add_argument(
    "-id",
    "--device_id",
    default=None,
    type=str,
    help="Device ID for GPU usage. If not provided, CPU is used.",
)
parser.add_argument(
    "-da",
    "--dataset",
    type=str,
    choices=["Bdataset", "Kdataset", "KGdataset", "KGdataset_tiny"],
    required=True,
    help="Dataset identifier for training. Options: 'Bdataset', 'Kdataset', 'KGdataset', or 'KGdataset_tiny'.",
)
parser.add_argument(
    "-sp",
    "--saved_path",
    type=str,
    default="result",
    help="Directory path to save training results.",
)
parser.add_argument(
    "-se",
    "--seed",
    default=42,
    type=int,
    help="Global random seed for reproducibility.",
)

# Training arguments
parser.add_argument(
    "-fo",
    "--nfold",
    default=10,
    type=int,
    help="Number of folds for K-fold cross-validation.",
)
parser.add_argument(
    "-ep",
    "--epoch",
    default=1000,
    type=int,
    help="Number of epochs for model training.",
)
parser.add_argument(
    "-lr",
    "--learning_rate",
    default=0.005,
    type=float,
    help="Learning rate for the optimizer.",
)
parser.add_argument(
    "-wd",
    "--weight_decay",
    default=0.0,
    type=float,
    help="Weight decay (L2 regularization) for the optimizer.",
)
parser.add_argument(
    "-pa",
    "--patience",
    default=100,
    type=int,
    help="Number of epochs with no improvement after which training will be stopped (early stopping).",
)

# Model hyperparameters
parser.add_argument(
    "-hf",
    "--hidden_feats",
    default=64,
    type=int,
    help="Dimension of hidden layers in the model.",
)
parser.add_argument(
    "-he",
    "--num_heads",
    default=5,
    type=int,
    help="Number of attention heads in the model.",
)
parser.add_argument(
    "-dp",
    "--dropout",
    default=0.0,
    type=float,
    help="Dropout rate to be applied in the model.",
)

# Parse the arguments and modify saved_path to include the seed for reproducibility
args = parser.parse_args()
args.saved_path = f"{args.saved_path}_{args.seed}"