"""
Base argument dataclass for Investigation experiments.

All project argument classes should inherit from :class:`ExampleArgs` to
ensure compatibility with ``ivst init``, ``ivst commit``, and the
:func:`investigation.decorator.main` decorator.
"""
from dataclasses import dataclass


@dataclass
class ExampleArgs:
    """Base dataclass for experiment hyperparameters.

    Subclass this in your project to define the full parameter space for your
    experiment.  All fields declared here are recognised by Investigation and will
    appear in the generated DoE JSON as ``shared`` parameters.

    Example:
        .. code-block:: python

            from dataclasses import dataclass
            from investigation.doeargs.args import ExampleArgs

            @dataclass
            class Args(ExampleArgs):
                name: str = "MyModel"
                learning_rate: float = 1e-3
                epochs: int = 100
    """

    #: Human-readable experiment name.  Used as a unique identifier across
    #: model variants in the DoE; each variant must have a distinct default.
    name: str = "Example"
    #: Base random seed.  Investigation generates ``num_seeds`` reproducible seeds
    #: from this value via :func:`numpy.random.default_rng`.
    seed: int = 1
    #: Number of CPU cores requested per Slurm task.
    num_cpus: int = 4
    #: Number of GPUs requested per Slurm task.
    num_gpus: int = 0
    #: Slurm wall-clock time limit in ``HH:MM:SS`` format.
    time: str = "01:00:00"
    #: Automatically requeue the job if it is preempted or hits the time limit.
    requeue: bool = False
    #: Compile the model with ``torch.compile`` before training (requires PyTorch ≥ 2.0).
    compile_model: bool = False


if __name__ == "__main__":
    print(ExampleArgs.mro())
    print(ExampleArgs.__annotations__)
    print(ExampleArgs.name)
    print(ExampleArgs.__dataclass_fields__)
    print({
        key: {
            "type": field.type,
            "default": field.default
        }
        for key, field in ExampleArgs.__dataclass_fields__.items()
    })
