import os

import json

from pathlib import Path

from pyhocon import ConfigTree, ConfigFactory

def get_conf(conf: ConfigTree, group: str = "", key: str = "", default=None):

    if group:

        key = ".".join([group, key])

    return conf.get(key, default)

class BaseConfig:

    GROUP = ""

    @classmethod

    def from_conf(cls, conf_path: Path):

        raise NotImplementedError("Must be implemented by subclasses")

    def save(self, output_dir: Path):

        assert output_dir.is_dir(), "output_dir has to be a directory."

        with open(str(output_dir.joinpath(self.__class__.__name__)) + ".json", "w") as file:

            json.dump(self.__dict__, file, indent=4)

    def to_dict(self):

        return self.__dict__

class ExperimentConfig(BaseConfig):

    GROUP = "experiment"

    def __init__(self, data_path: Path, output_dir: Path, n_evals: int, n_bootstraps: int, n_folds: int,

                 search_method: str):

        self.data_path = data_path

        self.output_dir = output_dir

        self.n_evals = n_evals

        self.n_folds = n_folds

        self.n_bootstraps = n_bootstraps

        self.search_method = search_method

class NNExperimentConfig(ExperimentConfig):

    def __init__(

            self, data_path: Path, output_dir: Path, n_evals: int, n_bootstraps: int, n_folds: int, search_method: str,

            batch_size: int, n_epochs: int, backend: str

    ):

        super().__init__(

            data_path=data_path, output_dir=output_dir, n_evals=n_evals, n_bootstraps=n_bootstraps, n_folds=n_folds,

            search_method=search_method

        )

        self.batch_size = batch_size

        self.n_epochs = n_epochs

        self.backend = backend

    @classmethod

    def from_conf(cls, conf_path):

        conf = ConfigFactory.parse_file(str(conf_path))

        data_path = Path(os.path.abspath(get_conf(conf, group=cls.GROUP, key="data_path")))

        if get_conf(conf, group=cls.GROUP, key="output_dir") is None:

            output_dir = data_path.parent.joinpath("output")

        else:

            output_dir = Path(get_conf(conf, group=cls.GROUP, key="output_dir"))

        return cls(

            data_path=data_path,

            output_dir=output_dir,

            batch_size=get_conf(conf, group=cls.GROUP, key="batch_size", default=16),

            n_epochs=get_conf(conf, group=cls.GROUP, key="n_epochs", default=100),

            n_evals=get_conf(conf, group=cls.GROUP, key="n_evals", default=50),

            n_bootstraps=get_conf(conf, group=cls.GROUP, key="n_bootstraps", default=100),

            n_folds=get_conf(conf, group=cls.GROUP, key="n_folds", default=6),

            search_method=get_conf(conf, group=cls.GROUP, key="search_method", default="particle swarm"),

            backend=get_conf(conf, group=cls.GROUP, key="backend", default="torch")

        )

class CoxExperimentConfig(ExperimentConfig):

    @classmethod

    def from_conf(cls, conf_path):

        conf = ConfigFactory.parse_file(str(conf_path))

        data_path = Path(os.path.abspath(get_conf(conf, group=cls.GROUP, key="data_path")))

        if get_conf(conf, group=cls.GROUP, key="output_dir") is None:

            output_dir = data_path.parent.joinpath("output")

        else:

            output_dir = Path(get_conf(conf, group=cls.GROUP, key="output_dir"))

        return cls(

            data_path=data_path,

            output_dir=output_dir,

            n_evals=get_conf(conf, group=cls.GROUP, key="n_evals", default=50),

            n_bootstraps=get_conf(conf, group=cls.GROUP, key="n_bootstraps", default=100),

            n_folds=get_conf(conf, group=cls.GROUP, key="n_folds", default=6),

            search_method=get_conf(conf, group=cls.GROUP, key="search_method", default="particle swarm")

        )

class HypersearchConfig(BaseConfig):

    GROUP = "hypersearch"

    def __init__(self, **kwargs):

        for key in kwargs:

            setattr(self, key, kwargs[key])

    @classmethod

    def from_conf(cls, conf_path: Path):

        conf = ConfigFactory.parse_file(str(conf_path))

        conf = getattr(conf, cls.GROUP)

        return cls(**conf)

class ModelConfig(BaseConfig):

    GROUP = "model"

    def __init__(self, **kwargs):

        for key in kwargs:

            setattr(self, key, kwargs[key])

    @classmethod

    def from_conf(cls, conf_path: Path):

        conf = ConfigFactory.parse_file(str(conf_path))

        conf = getattr(conf, cls.GROUP)

        return cls(**conf)