# Copyright 2019 The Feast Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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.
import pathlib
from collections import OrderedDict
from concurrent import futures
from datetime import datetime

import dataframes
import grpc
import pandas as pd
import pytest
import pytz
from google.protobuf import json_format

from feast.client import Client
from feast.core import CoreService_pb2_grpc as Core
from feast.entity import Entity
from feast.feature_set import (
    Feature,
    FeatureSet,
    FeatureSetRef,
    _make_tfx_schema_domain_info_inline,
)
from feast.value_type import ValueType
from feast_core_server import CoreServicer
from tensorflow_metadata.proto.v0 import schema_pb2

CORE_URL = "core.feast.local"
SERVING_URL = "serving.feast.local"


class TestFeatureSet:
    @pytest.fixture(scope="function")
    def server(self):
        server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
        Core.add_CoreServiceServicer_to_server(CoreServicer(), server)
        server.add_insecure_port("[::]:50051")
        server.start()
        yield server
        server.stop(0)

    @pytest.fixture
    def client(self, server):
        return Client(core_url="localhost:50051")

    def test_add_remove_features_success(self):
        fs = FeatureSet("my-feature-set")
        fs.add(Feature(name="my-feature-1", dtype=ValueType.INT64))
        fs.add(Feature(name="my-feature-2", dtype=ValueType.INT64))
        fs.drop(name="my-feature-1")
        assert len(fs.features) == 1 and fs.features[0].name == "my-feature-2"

    def test_remove_feature_failure(self):
        with pytest.raises(KeyError):
            fs = FeatureSet("my-feature-set")
            fs.drop(name="my-feature-1")

    def test_update_from_source_failure(self):
        with pytest.raises(Exception):
            df = pd.DataFrame()
            fs = FeatureSet("driver-feature-set")
            fs.infer_fields_from_df(df)

    @pytest.mark.parametrize(
        "dataframe,feature_count,entity_count,discard_unused_fields,features,entities",
        [
            (
                dataframes.GOOD,
                3,
                1,
                True,
                [],
                [Entity(name="entity_id", dtype=ValueType.INT64)],
            ),
            (
                dataframes.GOOD_FIVE_FEATURES,
                5,
                1,
                True,
                [],
                [Entity(name="entity_id", dtype=ValueType.INT64)],
            ),
            (
                dataframes.GOOD_FIVE_FEATURES,
                6,
                1,
                True,
                [Feature(name="feature_6", dtype=ValueType.INT64)],
                [Entity(name="entity_id", dtype=ValueType.INT64)],
            ),
            (
                dataframes.GOOD_FIVE_FEATURES_TWO_ENTITIES,
                5,
                2,
                True,
                [],
                [
                    Entity(name="entity_1_id", dtype=ValueType.INT64),
                    Entity(name="entity_2_id", dtype=ValueType.INT64),
                ],
            ),
            (
                dataframes.GOOD_FIVE_FEATURES_TWO_ENTITIES,
                6,
                3,
                False,
                [],
                [
                    Entity(name="entity_1_id", dtype=ValueType.INT64),
                    Entity(name="entity_2_id", dtype=ValueType.INT64),
                ],
            ),
            (
                dataframes.NO_FEATURES,
                0,
                1,
                True,
                [],
                [Entity(name="entity_id", dtype=ValueType.INT64)],
            ),
            (
                pd.DataFrame(
                    {
                        "datetime": [
                            datetime.utcnow().replace(tzinfo=pytz.utc) for _ in range(3)
                        ]
                    }
                ),
                0,
                0,
                True,
                [],
                [],
            ),
        ],
        ids=[
            "Test small dataframe update with hardcoded entity",
            "Test larger dataframe update with hardcoded entity",
            "Test larger dataframe update with hardcoded entity and feature",
            "Test larger dataframe update with two hardcoded entities and discarding of existing fields",
            "Test larger dataframe update with two hardcoded entities and retention of existing fields",
            "Test dataframe with no featuresdataframe",
            "Test empty dataframe",
        ],
    )
    def test_add_features_from_df_success(
        self,
        dataframe,
        feature_count,
        entity_count,
        discard_unused_fields,
        features,
        entities,
    ):
        my_feature_set = FeatureSet(
            name="my_feature_set",
            features=[Feature(name="dummy_f1", dtype=ValueType.INT64)],
            entities=[Entity(name="dummy_entity_1", dtype=ValueType.INT64)],
        )
        my_feature_set.infer_fields_from_df(
            dataframe,
            discard_unused_fields=discard_unused_fields,
            features=features,
            entities=entities,
        )
        assert len(my_feature_set.features) == feature_count
        assert len(my_feature_set.entities) == entity_count

    def test_import_tfx_schema(self):
        tests_folder = pathlib.Path(__file__).parent
        test_input_schema_json = open(
            tests_folder / "data" / "tensorflow_metadata" / "bikeshare_schema.json"
        ).read()
        test_input_schema = schema_pb2.Schema()
        json_format.Parse(test_input_schema_json, test_input_schema)

        feature_set = FeatureSet(
            name="bikeshare",
            entities=[Entity(name="station_id", dtype=ValueType.INT64)],
            features=[
                Feature(name="name", dtype=ValueType.STRING),
                Feature(name="status", dtype=ValueType.STRING),
                Feature(name="latitude", dtype=ValueType.FLOAT),
                Feature(name="longitude", dtype=ValueType.FLOAT),
                Feature(name="location", dtype=ValueType.STRING),
            ],
        )

        # Before update
        for entity in feature_set.entities:
            assert entity.presence is None
            assert entity.shape is None
        for feature in feature_set.features:
            assert feature.presence is None
            assert feature.shape is None
            assert feature.string_domain is None
            assert feature.float_domain is None
            assert feature.int_domain is None

        feature_set.import_tfx_schema(test_input_schema)

        # After update
        for feature in feature_set.features:
            assert feature.presence is not None
            assert feature.shape is not None
            if feature.name in ["location", "name", "status"]:
                assert feature.string_domain is not None
            elif feature.name in ["latitude", "longitude"]:
                assert feature.float_domain is not None
            elif feature.name in ["station_id"]:
                assert feature.int_domain is not None

    def test_export_tfx_schema(self):
        tests_folder = pathlib.Path(__file__).parent
        test_input_feature_set = FeatureSet.from_yaml(
            str(
                tests_folder
                / "data"
                / "tensorflow_metadata"
                / "bikeshare_feature_set.yaml"
            )
        )

        expected_schema_json = open(
            tests_folder / "data" / "tensorflow_metadata" / "bikeshare_schema.json"
        ).read()
        expected_schema = schema_pb2.Schema()
        json_format.Parse(expected_schema_json, expected_schema)
        _make_tfx_schema_domain_info_inline(expected_schema)

        actual_schema = test_input_feature_set.export_tfx_schema()

        assert len(actual_schema.feature) == len(expected_schema.feature)
        for actual, expected in zip(actual_schema.feature, expected_schema.feature):
            assert actual.SerializeToString() == expected.SerializeToString()

    def test_feature_set_import_export_yaml(self):

        test_feature_set = FeatureSet(
            name="bikeshare",
            entities=[Entity(name="station_id", dtype=ValueType.INT64)],
            features=[
                Feature(name="name", dtype=ValueType.STRING),
                Feature(name="longitude", dtype=ValueType.FLOAT),
                Feature(name="location", dtype=ValueType.STRING),
            ],
        )

        # Create a string YAML representation of the feature set
        string_yaml = test_feature_set.to_yaml()

        # Create a new feature set object from the YAML string
        actual_feature_set_from_string = FeatureSet.from_yaml(string_yaml)

        # Ensure equality is upheld to original feature set
        assert test_feature_set == actual_feature_set_from_string


def make_tfx_schema_domain_info_inline(schema):
    # Copy top-level domain info defined in the schema to inline definition.
    # One use case is in FeatureSet which does not have access to the top-level domain
    # info.
    domain_ref_to_string_domain = {d.name: d for d in schema.string_domain}
    domain_ref_to_float_domain = {d.name: d for d in schema.float_domain}
    domain_ref_to_int_domain = {d.name: d for d in schema.int_domain}

    for feature in schema.feature:
        domain_info_case = feature.WhichOneof("domain_info")
        if domain_info_case == "domain":
            domain_ref = feature.domain
            if domain_ref in domain_ref_to_string_domain:
                feature.string_domain.MergeFrom(domain_ref_to_string_domain[domain_ref])
            elif domain_ref in domain_ref_to_float_domain:
                feature.float_domain.MergeFrom(domain_ref_to_float_domain[domain_ref])
            elif domain_ref in domain_ref_to_int_domain:
                feature.int_domain.MergeFrom(domain_ref_to_int_domain[domain_ref])


def test_feature_set_class_contains_labels():
    fs = FeatureSet("my-feature-set", labels={"key1": "val1", "key2": "val2"})
    assert "key1" in fs.labels.keys() and fs.labels["key1"] == "val1"
    assert "key2" in fs.labels.keys() and fs.labels["key2"] == "val2"


def test_feature_class_contains_labels():
    fs = FeatureSet("my-feature-set", labels={"key1": "val1", "key2": "val2"})
    fs.add(
        Feature(
            name="my-feature-1",
            dtype=ValueType.INT64,
            labels={"feature_key1": "feature_val1"},
        )
    )
    assert "feature_key1" in fs.features[0].labels.keys()
    assert fs.features[0].labels["feature_key1"] == "feature_val1"


def test_feature_set_without_labels_empty_dict():
    fs = FeatureSet("my-feature-set")
    assert fs.labels == OrderedDict()
    assert len(fs.labels) == 0


def test_feature_without_labels_empty_dict():
    f = Feature("my feature", dtype=ValueType.INT64)
    assert f.labels == OrderedDict()
    assert len(f.labels) == 0


def test_set_label_feature_set():
    fs = FeatureSet("my-feature-set")
    fs.set_label("k1", "v1")
    assert fs.labels["k1"] == "v1"


def test_set_labels_overwrites_existing():
    fs = FeatureSet("my-feature-set")
    fs.set_label("k1", "v1")
    fs.set_label("k1", "v2")
    assert fs.labels["k1"] == "v2"


def test_remove_labels_empty_failure():
    fs = FeatureSet("my-feature-set")
    with pytest.raises(KeyError):
        fs.remove_label("key1")


def test_remove_labels_invalid_key_failure():
    fs = FeatureSet("my-feature-set")
    fs.set_label("k1", "v1")
    with pytest.raises(KeyError):
        fs.remove_label("key1")


def test_unequal_feature_based_on_labels():
    f1 = Feature(name="feature-1", dtype=ValueType.INT64, labels={"k1": "v1"})
    f2 = Feature(name="feature-1", dtype=ValueType.INT64, labels={"k1": "v1"})
    assert f1 == f2
    f3 = Feature(name="feature-1", dtype=ValueType.INT64)
    assert f1 != f3
    f4 = Feature(name="feature-1", dtype=ValueType.INT64, labels={"k1": "notv1"})
    assert f1 != f4


def test_unequal_feature_set_based_on_labels():
    fs1 = FeatureSet("my-feature-set")
    fs2 = FeatureSet("my-feature-set")
    assert fs1 == fs2
    fs1.set_label("k1", "v1")
    fs2.set_label("k1", "v1")
    assert fs1 == fs2
    fs2.set_label("k1", "unequal")
    assert not fs1 == fs2


def test_unequal_feature_set_other_has_no_labels():
    fs1 = FeatureSet("my-feature-set")
    fs2 = FeatureSet("my-feature-set")
    assert fs1 == fs2
    fs1.set_label("k1", "v1")
    assert not fs1 == fs2


def test_unequal_feature_other_has_no_labels():
    f1 = Feature(name="feature-1", dtype=ValueType.INT64, labels={"k1": "v1"})
    f2 = Feature(name="feature-1", dtype=ValueType.INT64)
    assert f1 != f2


class TestFeatureSetRef:
    def test_from_feature_set(self):
        feature_set = FeatureSet("test", "test")
        ref = FeatureSetRef.from_feature_set(feature_set)

        assert ref.name == "test"
        assert ref.project == "test"

    def test_str_ref(self):
        original_ref = FeatureSetRef(project="test", name="test")
        ref_str = repr(original_ref)
        parsed_ref = FeatureSetRef.from_str(ref_str)
        assert original_ref == parsed_ref
