import yaml
from pathlib import Path
import warnings

# https://stackoverflow.com/questions/9169025/how-can-i-add-a-python-tuple-to-a-yaml-file-using-pyyaml
class PrettySafeLoader(yaml.SafeLoader):
    def construct_python_tuple(self, node):
        return tuple(self.construct_sequence(node))

PrettySafeLoader.add_constructor(
    u'tag:yaml.org,2002:python/tuple',
    PrettySafeLoader.construct_python_tuple)


class Config:

    def __init__(self, benchmark, runtime_config_file, root_path, exp_name, wandb_username=None, wandb_project=None, checkpoint_path=None):

        with open(runtime_config_file, "r") as f:
            config_dict = yaml.load(f, Loader=PrettySafeLoader)

        self.evolutionary_search_config = EvolutionarySearchConfig(config_dict)
        self.runtime = RuntimeConfig(benchmark, config_dict, root_path, exp_name, wandb_username=wandb_username, wandb_project=wandb_project, checkpoint_path=checkpoint_path)
        
        # Get the benchmark config
        benchmark_config_path = Path(root_path, "configs", "benchmarks", f"{benchmark}.yaml")
        self.benchmark_config = BenchmarkConfig(benchmark_config_path)
        
        self.config_dict = config_dict

class RuntimeConfig():

    def __init__(self, benchmark, config_dict, root_path, exp_name, wandb_username=None, wandb_project=None, checkpoint_path=None):

        runtime_params = config_dict["runtime"]

        self.root_path = Path(root_path)

        for key, value in runtime_params.items():
            setattr(self, key, value)

        self.artifact_dir = Path(self.root_path, "artifacts", benchmark, exp_name)

        # Set all the derived attributes
        self.checkpoint_dir = Path(self.artifact_dir, "checkpoints") if checkpoint_path is None else checkpoint_path
        self.config_dir = Path(self.artifact_dir, "config")
        self.stats_dir = Path(self.artifact_dir, "stats")
        self.temp_dir = Path(self.artifact_dir, "temp")
        self.log_dir = Path(self.artifact_dir, "logs")
        self.supernet_config_dir = Path(self.artifact_dir, "supernet_configs")

        # WandB setup
        self.wandb = {
            "project": wandb_project,
            "entity": wandb_username,
            "group": exp_name
        }


class BenchmarkConfig:

    def __init__(self, task_config_file):
        
        with open(task_config_file, "r") as f:
            config_dict = yaml.load(f, Loader=PrettySafeLoader)

        for key, value in config_dict.items():
            setattr(self, key, value)


class EvolutionarySearchConfig:

    def __init__(self, config_dict):
        
        try:
            self.search_params = config_dict["evolutionary_search"]

            for key, value in self.search_params.items():
                setattr(self, key, value)
        except KeyError:
            pass
        except AttributeError:
            warnings.warn("Evolutionary search does not contain any parameters.")
