import random
import numpy as np
import torch
import logging
import hydra
from omegaconf import OmegaConf
from omegaconf import DictConfig as Config
from .scripts.datagen import run_training
from .scripts.eval import run_evaluation
from .scripts.preprocess_data import preprocess_data


_log = logging.getLogger(__name__)


@hydra.main(version_base=None, config_path='conf', config_name="datagen")
def run_main(cfg:Config) -> None:
    training_run_id = None
    gen_time = None

    if (cfg.run_preprocessing):
        preprocess_data(cfg)

    if (cfg.run_training):
        training_run_id, gen_time = run_training(cfg)
        _log.info(f"Training run ID: {training_run_id}")

    if (cfg.run_eval):
        strategy = cfg.datagen_method.params.sampling_strategy
        _log.info(f"Running evaluation for strategy '{strategy}'...")

        # Temporarily override sampling_strategy in the config
        cfg.datagen_method.params.sampling_strategy = strategy

        # Pass correct generation time
        if gen_time is not None:
            strategy_gen_time = gen_time[strategy] if isinstance(gen_time, dict) else None
        else:
            strategy_gen_time = None

        run_evaluation(cfg, training_run_id=training_run_id, gen_time=strategy_gen_time)


if __name__ == "__main__":
    run_main()
