import os
from pathlib import Path
from typing import Any, Dict

from determined.common import util

MASTER_SCHEME = "http"
MASTER_IP = "localhost"
MASTER_PORT = "8080"
DET_VERSION = None
DEFAULT_MAX_WAIT_SECS = 1800
MAX_TASK_SCHEDULED_SECS = 30
MAX_TRIAL_BUILD_SECS = 90


DEFAULT_TF1_CPU_IMAGE = "determinedai/environments:py-3.7-pytorch-1.7-tf-1.15-cpu-08f9c9b"
DEFAULT_TF2_CPU_IMAGE = (
    "determinedai/environments:py-3.8-pytorch-1.9-lightning-1.3-tf-2.4-cpu-08f9c9b"
)
DEFAULT_TF1_GPU_IMAGE = "determinedai/environments:cuda-10.2-pytorch-1.7-tf-1.15-gpu-08f9c9b"
DEFAULT_TF2_GPU_IMAGE = (
    "determinedai/environments:cuda-11.1-pytorch-1.9-lightning-1.3-tf-2.4-gpu-08f9c9b"
)

TF1_CPU_IMAGE = os.environ.get("TF1_CPU_IMAGE") or DEFAULT_TF1_CPU_IMAGE
TF2_CPU_IMAGE = os.environ.get("TF2_CPU_IMAGE") or DEFAULT_TF2_CPU_IMAGE
TF1_GPU_IMAGE = os.environ.get("TF1_GPU_IMAGE") or DEFAULT_TF1_GPU_IMAGE
TF2_GPU_IMAGE = os.environ.get("TF2_GPU_IMAGE") or DEFAULT_TF2_GPU_IMAGE
GPU_ENABLED = os.environ.get("DET_TEST_GPU_ENABLED", "1") not in ("0", "false")

PROJECT_ROOT_PATH = Path(__file__).resolve().parents[2]


def fixtures_path(path: str) -> str:
    return os.path.join(os.path.dirname(__file__), "fixtures", path)


def tutorials_path(path: str) -> str:
    return os.path.join(os.path.dirname(__file__), "../../examples/tutorials", path)


def cv_examples_path(path: str) -> str:
    return os.path.join(os.path.dirname(__file__), "../../examples/computer_vision", path)


def nlp_examples_path(path: str) -> str:
    return os.path.join(os.path.dirname(__file__), "../../examples/nlp", path)


def nas_examples_path(path: str) -> str:
    return os.path.join(os.path.dirname(__file__), "../../examples/nas", path)


def meta_learning_examples_path(path: str) -> str:
    return os.path.join(os.path.dirname(__file__), "../../examples/meta_learning", path)


def gan_examples_path(path: str) -> str:
    return os.path.join(os.path.dirname(__file__), "../../examples/gan", path)


def decision_trees_examples_path(path: str) -> str:
    return os.path.join(os.path.dirname(__file__), "../../examples/decision_trees", path)


def features_examples_path(path: str) -> str:
    return os.path.join(os.path.dirname(__file__), "../../examples/features", path)


def model_hub_examples_path(path: str) -> str:
    return os.path.join(os.path.dirname(__file__), "../../model_hub/examples", path)


def graphs_examples_path(path: str) -> str:
    return os.path.join(os.path.dirname(__file__), "../../examples/graphs", path)


def load_config(config_path: str) -> Any:
    with open(config_path) as f:
        config = util.safe_load_yaml_with_exceptions(f)
    return config


def make_master_url(suffix: str = "") -> str:
    return "{}://{}:{}/{}".format(MASTER_SCHEME, MASTER_IP, MASTER_PORT, suffix)


def set_global_batch_size(config: Dict[Any, Any], batch_size: int) -> Dict[Any, Any]:
    config = config.copy()
    config["hyperparameters"]["global_batch_size"] = batch_size
    return config


def set_slots_per_trial(config: Dict[Any, Any], slots: int) -> Dict[Any, Any]:
    config = config.copy()
    config.setdefault("resources", {})
    config["resources"]["slots_per_trial"] = slots
    return config


def set_max_length(config: Dict[Any, Any], max_length: Dict[str, int]) -> Dict[Any, Any]:
    config = config.copy()
    config["searcher"]["max_length"] = max_length
    return config


def set_min_validation_period(
    config: Dict[Any, Any], min_validation_period: Dict[str, int]
) -> Dict[Any, Any]:
    config = config.copy()
    config["min_validation_period"] = min_validation_period
    return config


def set_min_checkpoint_period(
    config: Dict[Any, Any], min_checkpoint_period: Dict[str, int]
) -> Dict[Any, Any]:
    config = config.copy()
    config["min_checkpoint_period"] = min_checkpoint_period
    return config


def set_aggregation_frequency(config: Dict[Any, Any], aggregation_frequency: int) -> Dict[Any, Any]:
    config = config.copy()
    config.setdefault("optimizations", {})
    config["optimizations"]["aggregation_frequency"] = aggregation_frequency
    return config


def set_tensor_auto_tuning(config: Dict[Any, Any], auto_tune: bool) -> Dict[Any, Any]:
    config = config.copy()
    config.setdefault("optimizations", {})
    config["optimizations"]["auto_tune_tensor_fusion"] = auto_tune
    return config


def set_image(config: Dict[Any, Any], cpu_image: str, gpu_image: str) -> Dict[Any, Any]:
    config = config.copy()
    config.setdefault("environment", {})
    config["environment"]["image"] = {"cpu": cpu_image, "gpu": gpu_image}
    return config


def set_tf1_image(config: Dict[Any, Any]) -> Dict[Any, Any]:
    return set_image(config, TF1_CPU_IMAGE, TF1_GPU_IMAGE)


def set_tf2_image(config: Dict[Any, Any]) -> Dict[Any, Any]:
    return set_image(config, TF2_CPU_IMAGE, TF2_GPU_IMAGE)


def set_shared_fs_data_layer(config: Dict[Any, Any]) -> Dict[Any, Any]:
    config = config.copy()
    config["data_layer"] = {}
    config["data_layer"]["type"] = "shared_fs"
    return config


def set_s3_data_layer(config: Dict[Any, Any]) -> Dict[Any, Any]:
    config = config.copy()
    config["data_layer"] = {}
    config["data_layer"]["type"] = "s3"
    config["data_layer"]["bucket"] = "yogadl-test"
    config["data_layer"]["bucket_directory_path"] = "determined_integration_tests"
    return config


def set_random_seed(config: Dict[Any, Any], seed: int) -> Dict[Any, Any]:
    config = config.copy()
    config.setdefault("reproducibility", {})
    config["reproducibility"]["experiment_seed"] = seed
    return config


def set_hparam(config: Dict[Any, Any], name: str, value: Any) -> Dict[Any, Any]:
    config = config.copy()
    config["hyperparameters"][name] = {"type": "const", "val": value}
    return config


def set_perform_initial_validation(config: Dict[Any, Any], init_val: bool) -> Dict[Any, Any]:
    config = config.copy()
    config["perform_initial_validation"] = init_val
    return config


def set_pod_spec(config: Dict[Any, Any], pod_spec: Dict[Any, Any]) -> Dict[Any, Any]:
    config = config.copy()
    config.setdefault("environment", {})
    config["environment"]["pod_spec"] = pod_spec
    return config


def set_profiling_enabled(config: Dict[Any, Any]) -> Dict[Any, Any]:
    config = config.copy()
    config.setdefault("profiling", {})
    config["profiling"]["enabled"] = True
    return config


def set_entrypoint(config: Dict[Any, Any], entrypoint: str) -> Dict[Any, Any]:
    config = config.copy()
    config["entrypoint"] = entrypoint
    return config

[((703, 734), 'os.environ.get', 'os.environ.get', (['"""TF1_CPU_IMAGE"""'], {}), "('TF1_CPU_IMAGE')\n", (717, 734), False, 'import os\n'), ((776, 807), 'os.environ.get', 'os.environ.get', (['"""TF2_CPU_IMAGE"""'], {}), "('TF2_CPU_IMAGE')\n", (790, 807), False, 'import os\n'), ((849, 880), 'os.environ.get', 'os.environ.get', (['"""TF1_GPU_IMAGE"""'], {}), "('TF1_GPU_IMAGE')\n", (863, 880), False, 'import os\n'), ((922, 953), 'os.environ.get', 'os.environ.get', (['"""TF2_GPU_IMAGE"""'], {}), "('TF2_GPU_IMAGE')\n", (936, 953), False, 'import os\n'), ((993, 1036), 'os.environ.get', 'os.environ.get', (['"""DET_TEST_GPU_ENABLED"""', '"""1"""'], {}), "('DET_TEST_GPU_ENABLED', '1')\n", (1007, 1036), False, 'import os\n'), ((1179, 1204), 'os.path.dirname', 'os.path.dirname', (['__file__'], {}), '(__file__)\n', (1194, 1204), False, 'import os\n'), ((1288, 1313), 'os.path.dirname', 'os.path.dirname', (['__file__'], {}), '(__file__)\n', (1303, 1313), False, 'import os\n'), ((1415, 1440), 'os.path.dirname', 'os.path.dirname', (['__file__'], {}), '(__file__)\n', (1430, 1440), False, 'import os\n'), ((1549, 1574), 'os.path.dirname', 'os.path.dirname', (['__file__'], {}), '(__file__)\n', (1564, 1574), False, 'import os\n'), ((1671, 1696), 'os.path.dirname', 'os.path.dirname', (['__file__'], {}), '(__file__)\n', (1686, 1696), False, 'import os\n'), ((1803, 1828), 'os.path.dirname', 'os.path.dirname', (['__file__'], {}), '(__file__)\n', (1818, 1828), False, 'import os\n'), ((1935, 1960), 'os.path.dirname', 'os.path.dirname', (['__file__'], {}), '(__file__)\n', (1950, 1960), False, 'import os\n'), ((2068, 2093), 'os.path.dirname', 'os.path.dirname', (['__file__'], {}), '(__file__)\n', (2083, 2093), False, 'import os\n'), ((2206, 2231), 'os.path.dirname', 'os.path.dirname', (['__file__'], {}), '(__file__)\n', (2221, 2231), False, 'import os\n'), ((2339, 2364), 'os.path.dirname', 'os.path.dirname', (['__file__'], {}), '(__file__)\n', (2354, 2364), False, 'import os\n'), ((2470, 2495), 'os.path.dirname', 'os.path.dirname', (['__file__'], {}), '(__file__)\n', (2485, 2495), False, 'import os\n'), ((2622, 2660), 'determined.common.util.safe_load_yaml_with_exceptions', 'util.safe_load_yaml_with_exceptions', (['f'], {}), '(f)\n', (2657, 2660), False, 'from determined.common import util\n'), ((1080, 1094), 'pathlib.Path', 'Path', (['__file__'], {}), '(__file__)\n', (1084, 1094), False, 'from pathlib import Path\n')]