CREATE SCHEMA offline;

CREATE TABLE offline.features (
  "ENTITY_ID"         VARCHAR,
  "EVENT_TIMESTAMP" TIMESTAMP,
  "ENTITY_FLOAT"      FLOAT,
);

INSERT INTO offline.features
SELECT *
FROM (
  VALUES ('11111111', '2024-01-01 13:00:00' :: TIMESTAMP, 1.1),
         ('11111111', '2024-01-01 14:00:00' :: TIMESTAMP, 1.11),
         ('11111111', '2024-01-01 15:00:00' :: TIMESTAMP, 1.111),
         ('22222222', '2024-01-01 13:00:00' :: TIMESTAMP, 2.2),
         ('22222222', '2024-01-01 14:00:00' :: TIMESTAMP, 2.22),
         ('33333333', '2024-01-01 13:00:00' :: TIMESTAMP, 3.3),
         ('44444444', '2024-01-02 22:00:00' :: TIMESTAMP, 4.4)
  )

from datetime import timedelta
from typing import Any

import pandas as pd
from feast import (
    Entity,
    FeatureService,
    FeatureStore,
    FeatureView,
    Field,
    RequestSource,
    ValueType,
)
from feast.infra.offline_stores.contrib.postgres_offline_store.postgres_source import (
    PostgreSQLSource as PostgresSource,
)
from feast.on_demand_feature_view import on_demand_feature_view
from feast.types import Float32, Float64

feature_store = FeatureStore()

features_entity = Entity(
    name="entity_id",
    join_keys=["ENTITY_ID"],
    value_type=ValueType.STRING,
)

features_source = PostgresSource(
    name="features",
    timestamp_field="EVENT_TIMESTAMP",
    table="offline.features",
)

features_feature_view = FeatureView(
    name="features_feature_view",
    entities=[features_entity],
    ttl=timedelta(days=0),
    schema=[Field(name="ENTITY_FLOAT", dtype=Float32)],
    online=True,
    source=features_source,
)

request_source = RequestSource(
    name="request_feature",
    schema=[Field(name="REQUEST_FLOAT", dtype=Float32)],
)


@on_demand_feature_view(
    sources=[features_feature_view, request_source],
    schema=[
        Field(name="ENTITY_FLOAT_TRANSFORMED_PANDAS", dtype=Float64),
        Field(name="ENTITY_FLOAT_PLUS_REQUEST_SOURCE", dtype=Float64),
    ],
    mode="pandas",
)
def odfv_pandas(input: pd.DataFrame) -> pd.DataFrame:
    output = pd.DataFrame()
    output["ENTITY_FLOAT_TRANSFORMED_PANDAS"] = input["ENTITY_FLOAT"] * 2
    output["ENTITY_FLOAT_PLUS_REQUEST_SOURCE"] = (
        input["ENTITY_FLOAT"] * input["REQUEST_FLOAT"]
    )
    return output


@on_demand_feature_view(
    sources=[features_feature_view, request_source],
    schema=[Field(name="ENTITY_FLOAT_TRANSFORMED_PYTHON", dtype=Float64)],
    mode="python",
)
def odfv_python(input: dict[str, Any]) -> dict[str, Any]:
    output = {}
    output["ENTITY_FLOAT_TRANSFORMED_PYTHON"] = [
        value * 2 if value is not None else None for value in input["ENTITY_FLOAT"]
    ]

    output["ENTITY_FLOAT_PLUS_REQUEST_SOURCE_PYTHON"] = [
        (e + r) if e is not None and r is not None else None
        for e, r in zip(input["ENTITY_FLOAT"], input["REQUEST_FLOAT"])
    ]
    return output


features_feature_service_pandas = FeatureService(
    name="features_feature_service_pandas",
    features=[features_feature_view, odfv_pandas],
)

features_feature_service_python = FeatureService(
    name="features_feature_service_python",
    features=[features_feature_view, odfv_python],
)

feature_store.apply(
    [
        features_entity,
        features_source,
        features_feature_view,
        odfv_pandas,
        odfv_python,
        features_feature_service_pandas,
        features_feature_service_python,
    ]
)

from datetime import datetime

from feast import FeatureStore

feature_store = FeatureStore()
feature_store.materialize(
    start_date=datetime(1900, 1, 1),
    end_date=datetime(9999, 1, 1),
    feature_views=["features_feature_view"],
)

import pandas as pd
from feast import FeatureStore

feature_store = FeatureStore()
feature_service_pandas = feature_store.get_feature_service(
    name="features_feature_service_pandas"
)
feature_service_python = feature_store.get_feature_service(
    name="features_feature_service_python"
)

entity_rows = [
    # This entity ID is not in the offline or online store
    {"ENTITY_ID": "1", "REQUEST_FLOAT": 1.0},
]
entity_df = pd.DataFrame(entity_rows)
entity_df["event_timestamp"] = pd.to_datetime("now", utc=True)

# This works.
print("offline with pandas")
offline_features = feature_store.get_historical_features(
    entity_df=entity_df,
    features=feature_service_pandas,
).to_df()
print(list(offline_features.to_dict().keys()))

## This doesn't work, raises the error
# print("online with pandas")
# online_features = feature_store.get_online_features(
#     entity_rows=entity_rows,
#     features=feature_service_pandas,
# ).to_dict()
# print(list(online_features.keys()))

## This doesn't work, raises the error
# print("online with python")
# online_features = feature_store.get_online_features(
#     entity_rows=entity_rows,
#     features=feature_service_python,
# ).to_dict()
# print(list(online_features.keys()))