#!/usr/bin/python3
"""
MAIN
====

+ Register

+ Setup
    - Seed
    - Device
    - Backends

+ Load
    - Dataset
    - Model

+ Main
    - CLI/Config
    - Logging
    - Setup
    - Load
    - Train
    - Load
    - Test

"""

import os
import time

import hydra
from omegaconf import OmegaConf
import wandb

import neurowave.data
from neurowave.registry import DATASET_REGISTRY

print('test')


# ========
# SETUP
# ========

# ========
# LOAD
# ========

# ========
# MAIN
# ========
@hydra.main(version_base=None, config_path='cfg', config_name='config')
def main(cfg):
    # Print config
    print(OmegaConf.to_yaml(cfg))

    # Initialize wandb
    if cfg.wandb.enabled:
        os.makedirs(cfg.wandb.dir, exist_ok=True)
        wandb.init(
            project=cfg.wandb.project,
            entity=cfg.wandb.entity,
            mode=cfg.wandb.mode,
            config=OmegaConf.to_container(cfg, resolve=True),
            name=cfg.wandb.name,
            dir=cfg.wandb.dir,
        )

    # Load dataset
    dataset = DATASET_REGISTRY[cfg.dataset.name](cfg.dataset)
    x, label = dataset.__getitem__(0)  # Load first item for demonstration
    print(f"Loaded item: {x.shape}, Label: {label.shape}")

    print(f"Dataset: {dataset.task}, Size: {len(dataset)}")

if __name__ == "__main__":
    start_time = time.time()
    main()
    end_time = time.time()
    print(f"Execution Time: {end_time - start_time:.2f} seconds")

