os

+""" general utility functions"""
+import argparse
+import importlib
+import json
+import logging
+import os
+import random
+import re
+import shutil
+import sys
+import typing
+from argparse import ArgumentParser
+from collections.abc import MutableMapping
+
+import numpy
+import torch
+from box import Box
+
+logger = logging.getLogger()
+
+
+def listorstr(inp):
+    if len(inp) == 1:
+        return try_cast(inp[0])
+
+    for i, val in enumerate(inp):
+        inp[i] = try_cast(val)
+    return inp
+
+
+def try_cast(text):
+    """ try to cast to int or float if possible, else return the text itself"""
+    result = try_int(text, None)
+    if result is not None:
+        return result
+
+    result = try_float(text, None)
+    if result is not None:
+        return result
+
+    return text
+
+
+def try_float(text, default: typing.Optional[int] = 0.0):
+    result = default
+    try:
+        result = float(text)
+    except Exception as _:
+        pass
+    return result
+
+
+def try_int(text, default: typing.Optional[int] = 0):
+    result = default
+    try:
+        result = int(text)
+    except Exception as _:
+        pass
+    return result
+
+
+def parse_args(parser: ArgumentParser) -> Box:
+    # get defaults
+    defaults = {}
+    # taken from parser_known_args code
+    # add any action defaults that aren't present
+    for action in parser._actions:
+        if action.dest is not argparse.SUPPRESS:
+            if action.default is not argparse.SUPPRESS:
+                defaults[action.dest] = action.default
+
+    # add any parser defaults that aren't present
+    for dest in parser._defaults:
+        defaults[dest] = parser._defaults[dest]
+
+    # check if there is config & read config
+    args = parser.parse_args()
+    if vars(args).get("config") is not None:
+        # load a .py config
+        configFile = args.config
+        spec = importlib.util.spec_from_file_location("config", configFile)
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        config = module.config
+        # merge config and override defaults
+        defaults.update({k: v for k, v in config.items()})
+
+    # override defaults with command line params
+    # this will get rid of defaults and only read command line args
+    parser._defaults = {}
+    parser._actions = {}
+    args = parser.parse_args()
+    defaults.update({k: v for k, v in vars(args).items()})
+
+    return boxify_dict(defaults)
+
+
+def boxify_dict(config):
+    """
+  this takes a flat dictionary and break it into sub-dictionaries based on "." seperation
+    a = {"model.a":  1, "model.b" : 2,  "alpha" : 3} will return Box({"model" : {"a" :1,
+    "b" : 2}, alpha:3})
+    a = {"model.a":  1, "model.b" : 2,  "model" : 3} will throw error
+  """
+    new_config = {}
+    # iterate over keys and split on "."
+    for key in config:
+        if "." in key:
+            temp_config = new_config
+            for k in key.split(".")[:-1]:
+                # create non-existent keys as dictionary recursively
+                if temp_config.get(k) is None:
+                    temp_config[k] = {}
+                elif not isinstance(temp_config.get(k), dict):
+                    raise TypeError(f"Key '{k}' has values as well as child")
+                temp_config = temp_config[k]
+            temp_config[key.split(".")[-1]] = config[key]
+        else:
+            if new_config.get(key) is None:
+                new_config[key] = config[key]
+            else:
+                raise TypeError(f"Key '{key}' has values as well as child")
+
+    return Box(new_config)
+
+
+# https://stackoverflow.com/questions/6027558/flatten-nested-dictionaries-compressing-keys
+def flatten(d, parent_key='', sep='.'):
+    items = []
+    for k, v in d.items():
+        new_key = parent_key + sep + k if parent_key else k
+        if isinstance(v, MutableMapping):
+            items.extend(flatten(v, new_key, sep=sep).items())
+        else:
+            items.append((new_key, v))
+    return Box(dict(items))
+
+
+def str2bool(v: typing.Union[bool, str, int]) -> bool:
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ("yes", "true", "t", "y", "1", 1):
+        return True
+    if v.lower() in ("no", "false", "f", "n", "0", 0):
+        return False
+    raise TypeError("Boolean value expected.")
+
+
+def safe_isdir(dir_name):
+    return os.path.exists(dir_name) and os.path.isdir(dir_name)
+
+
+def safe_makedirs(dir_name):
+    try:
+        os.makedirs(dir_name)
+    except OSError as e:
+        print(e)
+
+
+def jsonize(x: object) -> typing.Union[str, dict]:
+    try:
+        temp = json.dumps(x)
+        return temp
+    except Exception as e:
+        return {}
+
+
+def copy_code(folder_to_copy, out_folder, replace=False):
+    logger.info(f"copying {folder_to_copy} to {out_folder}")
+
+    if os.path.exists(out_folder):
+        if not os.path.isdir(out_folder):
+            logger.error(f"{out_folder} is not a directory")
+            sys.exit()
+        else:
+            logger.info(f"Not deleting existing result folder: {out_folder}")
+    else:
+        os.makedirs(out_folder)
+
+    # replace / with _
+    folder_name = f'{out_folder}/{re.sub("/", "_", folder_to_copy)}'
+
+    # create a new copy if something already exists
+    if not replace:
+        i = 1
+        temp = folder_name
+        while os.path.exists(temp):
+            temp = f"{folder_name}_{i}"
+            i += 1
+        folder_name = temp
+    else:
+        if os.path.exists(folder_name):
+            if os.path.isdir(folder_name):
+                shutil.rmtree(folder_name)
+            else:
+                raise FileExistsError("There is a file with same name as folder")
+
+    logger.info(f"Copying {folder_to_copy} to {folder_name}")
+    shutil.copytree(folder_to_copy, folder_name)
+
+
+def get_state_params(wandb_use, run_id, result_folder, statefile):
+    """This searches for model and run id in result folder
+        The logic is as follows
+
+    if we are not given run_id there are four cases:
+        - we want to restart the wandb run but too lazy to look up run-id or/and statefile
+        - we want a new wandb fresh run
+        - we are not using wandb at all and need to restart
+        - we are not using wandb and need a fresh run
+
+    Case 1/3:
+        - If we want to restart the run, we expect the result_folder name to end with
+        /run_<numeric>.
+        - In this case, if we are using wandb then we need to go inside wandb folder, list all
+        directory and pick up run id and (or) statefile
+        - If we are not using wandb we just look for model inside the run_<numeric> folder and
+        return statefile, run id as none
+
+    case 2/4:
+        if not 1/3, it is case 2/4
+
+    This is expected to be a fail safe script. i.e any of run_id or statefile may not be specified
+    and relies on whims of the user _-_
+    """
+    # if not resume get run number and create result_folder/run_{run_num}
+    # if someone is resuming we expect them to give the exact folder name upto run num.
+
+    # this part of code searches for run_id i.e will work only if we are using wandb
+    if run_id is None:
+        # if result folder if of type folder/run_<num>, then search for current checkpoint and
+        # run-id else we will just create a new run with run_<num+1>
+
+        regex = r"^.*/?run_[0-9]+/?$"
+        if re.match(regex, result_folder):
+
+            # search for checkpoint and run-id if using wandb
+            if wandb_use:
+                # search in wandb folder if it exists else we want a new run
+                if os.path.exists(f"{result_folder}/wandb/"):
+                    # case 1
+                    for folder in sorted(os.listdir(f"{result_folder}/wandb/"), reverse=True):
+                        # assume run_<##> will have only single run,
+                        # also no other crap in this folder
+                        if os.path.exists(f"{result_folder}/wandb/{folder}/current_model.pt"):
+                            run_id = folder.split("-")[-1]
+                            logger.info(f"using run id {run_id}")
+                        # we are done break out of for loop
+                        break
+            else:
+                # case 3
+                # if not using wandb search within run_<num> directory
+                logger.info(f"not using wandb")
+                if os.path.exists(f"{result_folder}/current_model.pt"):
+                    statefile = f"{result_folder}/current_model.pt"
+                    logger.info(f"using statefile {statefile}")
+                else:
+                    # just start a new run
+                    pass
+        else:
+            # trailing is not run_<num>; that means user wants a new fresh run
+            # so we give a fresh run and create a new folder
+            # case 2/4
+            last_run_num = max(
+                [0] + [try_int(i[-4:]) for i in os.listdir(result_folder)]) + 1
+            result_folder = f"{result_folder}/run_{last_run_num:04d}"
+            logger.info(f"Creating new run with {result_folder}")
+            safe_makedirs(result_folder)
+
+    # search for last checkpoint in case --statefile is none and we are resuming
+    if run_id is not None and statefile is None:
+        folders = sorted(os.listdir(f"{result_folder}/wandb"), reverse=True)
+        for folder in folders:
+            if run_id in folder:
+                # check for current_model.pt
+                if os.path.exists(f"{result_folder}/wandb/{folder}/current_model.pt"):
+                    statefile = f"{result_folder}/wandb/{folder}/current_model.pt"
+                    logger.info(f"Using state file {statefile} and run id {run_id}")
+                    break
+        if statefile is None:
+            raise Exception("Did not find statefile, exiting!!")
+    return statefile, run_id, result_folder
+
+
+if __name__ == "__main__":
+    # test boxify_dict
+    a = {"model.a": 1, "m   odel.b": 2, "alpha": 3}
+    print(boxify_dict(a))
+
+    try:
+        a = {"model.a": 1, "model.b": 2, "model": 3}
+        print(boxify_dict(a))
+    except Exception as e:
+        print(e)
+
+    try:
+        a = {"model": 4, "model.a": 1, "model.b": 2, "model": 3}
+        print(boxify_dict(a))
+    except Exception as e:
+        print(e)
+
+    try:
+        a = {"model.a": 1, "model": 4, "model.b": 2, "model": 3}
+        print(boxify_dict(a))
+    except Exception as e:
+        print(e)
+
+    try:
+        a = {"model": {"attr1": 1, "attr2": {"attr_attr_3": 3}}, "train": 10}
+        print(flatten(a))
+    except Exception as e:
+        print(e)
+
+
+def set_seed(seed):
+    if isinstance(seed, list):
+        torch_seed, numpy_seed, random_seed = seed
+    else:
+        torch_seed, numpy_seed, random_seed = seed, seed, seed
+
+    torch.manual_seed(torch_seed)
+    numpy.random.seed(numpy_seed)
+    random.seed(random_seed)