# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/ml.experiment.ipynb.

# %% auto 0
__all__ = ['make_experiment_dir', 'before_experiment', 'after_experiment', 'make_run_experiment_sweep', 'main']

# %% ../../nbs/ml.experiment.ipynb 3
import os
from pathlib import Path
import wandb
import json
from ..utils import NestedDict, flatten_dict

# %% ../../nbs/ml.experiment.ipynb 4
def make_experiment_dir(root=".experiments", name=None):
    if name is None:
        from bellek.utils import generate_time_id
        name = generate_time_id()
    experiment_dir = Path(root) / name
    experiment_dir.mkdir(parents=True, exist_ok=True)
    return experiment_dir

# %% ../../nbs/ml.experiment.ipynb 5
def before_experiment(wandb_run):
    config = NestedDict.from_flat_dict(wandb_run.config)

    # W&B variables
    config.set("wandb.id", wandb_run.id)
    os.environ["WANDB_PROJECT"] = wandb_run.project
    os.environ["WANDB_LOG_MODEL"] = "end"

    # Set random seed
    if seed := config.get("seed"):
        from fastai.torch_core import set_seed

        set_seed(seed)

    pp_config = config.at("metaconfig.preprocessing")
    
    # Resolve paths
    if pp_config.get("resolve_paths"):
        exclude_resolving_paths = pp_config.get("exclude_resolving_paths", [])
        for k, v in config.flat().items():
            if isinstance(k, str) and k.endswith("path") and k not in exclude_resolving_paths:
                config.set(k, str(Path(v).resolve()))

    # Resolve device
    if "device" not in config and pp_config.get("resolve_device"):
        import torch

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

    # Update W&B config
    wandb_run.config.update(flatten_dict(config), allow_val_change=True)

    # Save preprocessed config
    with open("./config.before.json", "w") as f:
        json.dump(config, f, indent=2)


def after_experiment(wandb_run):
    config = NestedDict.from_flat_dict(wandb_run.config)
    with open("./config.after.json", "w") as f:
        json.dump(config, f, indent=2)

# %% ../../nbs/ml.experiment.ipynb 6
def make_run_experiment_sweep(run_experiment, config_defaults):
    def func():
        wandb_params = config_defaults["wandb"]
        with wandb.init(config=flatten_dict(config_defaults), **wandb_params) as wandb_run:
            before_experiment(wandb_run)
            run_experiment(wandb_run)
            after_experiment(wandb_run)
    return func

def main(run_experiment, args):
    with open(args.cfg) as f:
        config = NestedDict(json.load(f))

    is_sweep = hasattr(args, "sweep_cfg") and args.sweep_cfg
    if is_sweep:
        with open(args.sweep_cfg) as f:
            sweep_config = json.load(f)
    else:
        sweep_config = {}

    run_experiment_sweep = make_run_experiment_sweep(run_experiment, config)
    if is_sweep:
        wandb_params = config["wandb"]
        count = sweep_config.pop("count") if "count" in sweep_config else None
        sweep_id = wandb.sweep(
            sweep_config,
            entity=wandb_params["entity"],
            project=wandb_params["project"],
        )
        wandb.agent(sweep_id, run_experiment_sweep, count=count)
    else:
        run_experiment_sweep()

