# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
from __future__ import absolute_import

from mock import Mock, patch
from packaging import version
import pytest

from sagemaker.tensorflow import TensorFlow

REGION = "us-west-2"

ENV_INPUT = {"env_key1": "env_val1", "env_key2": "env_val2", "env_key3": "env_val3"}


@pytest.fixture()
def sagemaker_session():
    return Mock(name="sagemaker_session", boto_region_name=REGION)


def _build_tf(sagemaker_session, **kwargs):
    return TensorFlow(
        sagemaker_session=sagemaker_session,
        entry_point="dummy.py",
        role="dummy-role",
        instance_count=1,
        instance_type="ml.c4.xlarge",
        **kwargs,
    )


@patch("sagemaker.fw_utils.python_deprecation_warning")
def test_estimator_py2_deprecation_warning(warning, sagemaker_session):
    estimator = _build_tf(sagemaker_session, framework_version="2.1.1", py_version="py2")

    assert estimator.py_version == "py2"
    warning.assert_called_with("tensorflow", "2.1.1")


def test_py2_version_deprecated(sagemaker_session):
    with pytest.raises(AttributeError) as e:
        _build_tf(sagemaker_session, framework_version="2.1.2", py_version="py2")

    msg = (
        "Python 2 containers are only available with 2.1.1 and lower versions. "
        "Please use a Python 3 container."
    )
    assert msg in str(e.value)


def test_py2_version_is_not_deprecated(sagemaker_session):
    estimator = _build_tf(sagemaker_session, framework_version="1.15.0", py_version="py2")
    assert estimator.py_version == "py2"
    estimator = _build_tf(sagemaker_session, framework_version="2.0.0", py_version="py2")
    assert estimator.py_version == "py2"


def test_framework_name(sagemaker_session):
    tf = _build_tf(sagemaker_session, framework_version="1.15.2", py_version="py3")
    assert tf._framework_name == "tensorflow"


def test_tf_add_environment_variables(sagemaker_session):
    tf = _build_tf(
        sagemaker_session,
        framework_version="1.15.2",
        py_version="py3",
        environment=ENV_INPUT,
    )
    assert tf.environment == ENV_INPUT


def test_tf_miss_environment_variables(sagemaker_session):
    tf = _build_tf(
        sagemaker_session,
        framework_version="1.15.2",
        py_version="py3",
        environment=None,
    )
    assert not tf.environment


def test_enable_sm_metrics(sagemaker_session):
    tf = _build_tf(
        sagemaker_session,
        framework_version="1.15.2",
        py_version="py3",
        enable_sagemaker_metrics=True,
    )
    assert tf.enable_sagemaker_metrics


def test_disable_sm_metrics(sagemaker_session):
    tf = _build_tf(
        sagemaker_session,
        framework_version="1.15.2",
        py_version="py3",
        enable_sagemaker_metrics=False,
    )
    assert not tf.enable_sagemaker_metrics


def test_disable_sm_metrics_if_fw_ver_is_less_than_1_15(
    sagemaker_session, tensorflow_training_version, tensorflow_training_py_version
):
    if version.Version(tensorflow_training_version) > version.Version("1.14"):
        pytest.skip("This test is for TF 1.14 and lower.")

    tf = _build_tf(
        sagemaker_session,
        framework_version=tensorflow_training_version,
        py_version=tensorflow_training_py_version,
        image_uri="old-image",
    )
    assert tf.enable_sagemaker_metrics is None


def test_enable_sm_metrics_if_fw_ver_is_at_least_1_15(
    sagemaker_session, tensorflow_training_version, tensorflow_training_py_version
):
    if version.Version(tensorflow_training_version) < version.Version("1.15"):
        pytest.skip("This test is for TF 1.15 and higher.")

    tf = _build_tf(
        sagemaker_session,
        framework_version=tensorflow_training_version,
        py_version=tensorflow_training_py_version,
    )
    assert tf.enable_sagemaker_metrics


def test_require_image_uri_if_fw_ver_is_less_than_1_11(
    sagemaker_session, tensorflow_training_version, tensorflow_training_py_version
):
    if version.Version(tensorflow_training_version) > version.Version("1.10"):
        pytest.skip("This test is for TF 1.10 and lower.")

    with pytest.raises(ValueError) as e:
        _build_tf(
            sagemaker_session,
            framework_version=tensorflow_training_version,
            py_version=tensorflow_training_py_version,
        )

    expected_msg = (
        "TF {version} supports only legacy mode. Please supply the image URI directly with "
        "'image_uri=520713654638.dkr.ecr.{region}.amazonaws.com/"
        "sagemaker-tensorflow:{version}-cpu-py2' and set 'model_dir=False'. If you are using any "
        "legacy parameters (training_steps, evaluation_steps, checkpoint_path, requirements_file), "
        "make sure to pass them directly as hyperparameters instead."
    ).format(version=tensorflow_training_version, region=REGION)

    assert expected_msg in str(e.value)

[(26, 'pytest.fixture', 'pytest.fixture', False, 'import pytest\n'), (42, 'mock.patch', 'patch', False, 'from mock import Mock, patch\n'), (28, 'mock.Mock', 'Mock', False, 'from mock import Mock, patch\n'), (32, 'sagemaker.tensorflow.TensorFlow', 'TensorFlow', False, 'from sagemaker.tensorflow import TensorFlow\n'), (51, 'pytest.raises', 'pytest.raises', False, 'import pytest\n'), (116, 'packaging.version.Version', 'version.Version', False, 'from packaging import version\n'), (116, 'packaging.version.Version', 'version.Version', False, 'from packaging import version\n'), (117, 'pytest.skip', 'pytest.skip', False, 'import pytest\n'), (131, 'packaging.version.Version', 'version.Version', False, 'from packaging import version\n'), (131, 'packaging.version.Version', 'version.Version', False, 'from packaging import version\n'), (132, 'pytest.skip', 'pytest.skip', False, 'import pytest\n'), (145, 'packaging.version.Version', 'version.Version', False, 'from packaging import version\n'), (145, 'packaging.version.Version', 'version.Version', False, 'from packaging import version\n'), (146, 'pytest.skip', 'pytest.skip', False, 'import pytest\n'), (148, 'pytest.raises', 'pytest.raises', False, 'import pytest\n')]