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

[((25, 16, 25, 47), 'os.environ.get', 'os.environ.get', ({(25, 31, 25, 46): '"""TF1_CPU_IMAGE"""'}, {}), "('TF1_CPU_IMAGE')", False, 'import os\n'), ((26, 16, 26, 47), 'os.environ.get', 'os.environ.get', ({(26, 31, 26, 46): '"""TF2_CPU_IMAGE"""'}, {}), "('TF2_CPU_IMAGE')", False, 'import os\n'), ((27, 16, 27, 47), 'os.environ.get', 'os.environ.get', ({(27, 31, 27, 46): '"""TF1_GPU_IMAGE"""'}, {}), "('TF1_GPU_IMAGE')", False, 'import os\n'), ((28, 16, 28, 47), 'os.environ.get', 'os.environ.get', ({(28, 31, 28, 46): '"""TF2_GPU_IMAGE"""'}, {}), "('TF2_GPU_IMAGE')", False, 'import os\n'), ((29, 14, 29, 57), 'os.environ.get', 'os.environ.get', ({(29, 29, 29, 51): '"""DET_TEST_GPU_ENABLED"""', (29, 53, 29, 56): '"""1"""'}, {}), "('DET_TEST_GPU_ENABLED', '1')", False, 'import os\n'), ((35, 24, 35, 49), 'os.path.dirname', 'os.path.dirname', ({(35, 40, 35, 48): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((39, 24, 39, 49), 'os.path.dirname', 'os.path.dirname', ({(39, 40, 39, 48): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((43, 24, 43, 49), 'os.path.dirname', 'os.path.dirname', ({(43, 40, 43, 48): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((47, 24, 47, 49), 'os.path.dirname', 'os.path.dirname', ({(47, 40, 47, 48): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((51, 24, 51, 49), 'os.path.dirname', 'os.path.dirname', ({(51, 40, 51, 48): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((55, 24, 55, 49), 'os.path.dirname', 'os.path.dirname', ({(55, 40, 55, 48): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((59, 24, 59, 49), 'os.path.dirname', 'os.path.dirname', ({(59, 40, 59, 48): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((63, 24, 63, 49), 'os.path.dirname', 'os.path.dirname', ({(63, 40, 63, 48): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((67, 24, 67, 49), 'os.path.dirname', 'os.path.dirname', ({(67, 40, 67, 48): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((71, 24, 71, 49), 'os.path.dirname', 'os.path.dirname', ({(71, 40, 71, 48): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((75, 24, 75, 49), 'os.path.dirname', 'os.path.dirname', ({(75, 40, 75, 48): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((80, 17, 80, 55), 'determined.common.util.safe_load_yaml_with_exceptions', 'util.safe_load_yaml_with_exceptions', ({(80, 53, 80, 54): 'f'}, {}), '(f)', False, 'from determined.common import util\n'), ((31, 20, 31, 34), 'pathlib.Path', 'Path', ({(31, 25, 31, 33): '__file__'}, {}), '(__file__)', False, 'from pathlib import Path\n')]