from path_learning.utils import ROOT_DIR

CONFIG_DIR = ROOT_DIR / "configs"

"""
Naming convention of new config files:
base_config_DATASET_TRANSFORM.json
"""


def write_config(config, **kwargs):
    """
    REPLACEMENTS:
    TRANSFORMTITLE -> Capitalized version of transform name
    TRANSFORM - lowercase version of transform name
    DATASET - Name of dataset to be substituted
    SLOPE - String of float of leaky-ReLU slope to use for neural network experiments
    :param config:
    :param kwargs:
    :return:
    """
    transform = kwargs["transform"]
    dataset = kwargs["dataset"]
    print(f"Using base {config}")

    # READ ALL LINES OF CONFIG FILE
    with open(str(CONFIG_DIR / config) + ".json", "r") as json_file:
        lines = json_file.readlines()
        print(lines)

    # EDIT LINES WHERE NECESSARY, WRITE NEW CONFIG FILE
    with open(str(CONFIG_DIR / config)+"_"+dataset+"_"+transform + ".json", 'w') as write_file:
        for line in lines:
            # EDITS - can be made more elaborate
            write_line = line.replace("BASETITLE", str(config))
            write_line = write_line.replace("TRANSFORMTITLE", transform.capitalize())
            write_line = write_line.replace("TRANSFORM", transform)
            write_line = write_line.replace("TRANSFORM", transform)
            write_line = write_line.replace("DATASET", dataset)
            write_line = write_line.replace('"SLOPE"', slope)
            # Writes edited line and goes to next line
            write_file.write(write_line)


if __name__ == "__main__":
    """
    EDIT THE FOLLOWING TO CREATE MORE CONFIG FILES FROM A BASE CONFIG FILE
        The base config file need to contain CODEWORDS described in the function "write_config" above
    """
    print(f"Using config directory {CONFIG_DIR}")

    # EDIT WHICH BASE-CONFIGS TO AMEND
    base_configs = ["N=2"]  # , "4_step", "6_step", "8_step", "10_step"]

    # EDIT WHICH DATASETS YOU WANT TO USE
    DATASETS = ["CIFAR10", "MNIST", "FashionMNIST"]

    # EDIT WHICH TRANSFORMS YOU WANT TO USE
    TRANSFORMS = ["blur", "contrast", "shift"]

    # EDIT WHICH SLOPE YOU WANT TO USE
    SLOPES = ['0.0']  # Note this way the "" are not written into the file

    for base_config in base_configs:
        for dataset in DATASETS:
            for transform in TRANSFORMS:
                for slope in SLOPES:
                    configuration = {"dataset": dataset,
                                     "transform": transform,
                                     "slope": slope}
                    write_config(base_config, **configuration)
