import json

from PhenPred.vae import data_folder, plot_folder

from PhenPred.vae.Losses import CLinesLosses

class Hypers:

    @classmethod

    def read_json(cls, json_file):

        with open(json_file, "r") as f:

            hypers = json.load(f)

        return hypers

    @classmethod

    def read_hyperparameters(

        cls, hypers_json=None, parse_torch_functions=True, timestamp=None

    ):

        if timestamp is not None:

            hypers_json = f"{plot_folder}/files/{timestamp}_hyperparameters.json"

        elif hypers_json is None:

            hypers_json = f"{plot_folder}/files/hyperparameters.json"

        hypers = cls.read_json(hypers_json)

        if timestamp is not None:

            hypers["load_run"] = timestamp

        if "model" not in hypers:

            hypers["model"] = "MOSA"

        if "standardize" not in hypers:

            hypers["standardize"] = False

        if "w_rec" not in hypers:

            hypers["w_rec"] = 1

        if "w_gauss" not in hypers:

            hypers["w_gauss"] = 0.01

        if "w_cat" not in hypers:

            hypers["w_cat"] = 0.01

        if hypers["view_loss_weights"] is None:

            hypers["view_loss_weights"] = [1.0] * len(hypers["views"])

        if hypers["use_conditionals"] is None:

            hypers["use_conditionals"] = True

        if timestamp is None:  # full path is already stored in previous json config

            hypers["datasets"] = {

                k: f"{data_folder}/{v}" for k, v in hypers["datasets"].items()

            }

        print(f"# ---- Hyperparameters")

        print(json.dumps(hypers, indent=4, sort_keys=True))

        if parse_torch_functions:

            hypers = cls.parse_torch_functions(hypers)

        return hypers

    @classmethod

    def parse_torch_functions(cls, hypers):

        if type(hypers["activation_function"]) == str:

            hypers["activation_function"] = CLinesLosses.activation_function(

                hypers["activation_function"]

            )

        if type(hypers["reconstruction_loss"]) == str:

            hypers["reconstruction_loss"] = CLinesLosses.reconstruction_loss_method(

                hypers["reconstruction_loss"]

            )

        if type(hypers["hidden_dims"]) == str:

            hypers["hidden_dims"] = [float(l) for l in hypers["hidden_dims"].split(",")]

        return hypers