%pip install tencent-wedata-feature-engineering

%pip list | grep tencent-wedata-feature-engineering

from wedata.feature_store.client import FeatureStoreClient
# Build a feature engineering client instance
# create FeatureStoreClient
client = FeatureStoreClient(spark, cloud_secret_id=cloud_secret_id, cloud_secret_key=cloud_secret_key)

from wedata.feature_store.constants.engine_types import EngineTypes
# Create user feature table--EMR on hive
# create user's feature table--EMR on hive
feature_table = client.create_table(
    name=table_name,  # Table name
    primary_keys=["wine_id"],      # Primary key
    df=wine_features_df,           # data
    engine_type=EngineTypes.HIVE_ENGINE,
    data_source_name=data_source_name,
    timestamp_key="event_timestamp",
    tags={ # Business tag
        "purpose": "demo", 
        "create_by": "wedata"
    }
)

# Create user feature table--DLC on iceberg (catalog not enabled)
# create user's feature table--DLC on iceberg(disable catalog)
feature_table = client.create_table(
    name=table_name,  # Table name
    database_name=database_name,
    primary_keys=["wine_id"],      # Primary key
    df=wine_features_df,           # data
    engine_type=EngineTypes.ICEBERG_ENGINE,
    data_source_name=data_source_name,
    timestamp_key="event_timestamp",
    tags={ # Business tag
        "purpose": "demo", 
        "create_by": "wedata"
    }
)

# Create user feature table--DLC on iceberg (catalog enabled)
# create user's feature table--DLC on iceberg(enable catalog)
feature_table = client.create_table(
    name=table_name,  # Table name
    primary_keys=["wine_id"],      # Primary key
    df=wine_features_df,           # data
    engine_type=EngineTypes.ICEBERG_ENGINE,
    data_source_name=data_source_name,
    timestamp_key="event_timestamp",
    tags={ # Business tag
        "purpose": "demo", 
        "create_by": "wedata"
    },
    catalog_name=catalog_name
)

from wedata.feature_store.constants.engine_types import EngineTypes
# Register feature table--EMR on hive
# register feature table--EMR on hive
client.register_table(database_name=database_name, name=register_table_name, timestamp_key="event_timestamp",engine_type=EngineTypes.HIVE_ENGINE, data_source_name=data_source_name, primary_keys=["wine_id",])

# Register feature table--DLC on iceberg (catalog not enabled)
# register feature table--DLC on iceberg(disable catalog)
client.register_table(database_name=database_name, name=register_table_name, timestamp_key="event_timestamp",engine_type=EngineTypes.ICEBERG_ENGINE, data_source_name=data_source_name, primary_keys=["wine_id",])

# register feature table--DLC on iceberg(enable catalog)
# register feature table--DLC on iceberg(enable catalog)
client.register_table(database_name=database_name, name=register_table_name, timestamp_key="event_timestamp",engine_type=EngineTypes.HIVE_ENGINE, data_source_name=data_source_name, primary_keys=["wine_id",], catalog_name=catalog_name)

# Read offline feature table--EMR on hive--&--DLC on iceberg(catalog not enabled)
# read offline feature table--EMR on hive--&--DLC on iceberg(disable catalog)
get_df = client.read_table(name=table_name, database_name=database_name)
get_df.show(30)

# Read online feature table--EMR on hive--&--DLC on iceberg(catalog not enabled)
# read online feature table--EMR on hive--&--DLC on iceberg(disable catalog)
online_config = RedisStoreConfig(host=redis_host, port=redis_password, db=redis_db, password=redis_password)
primary_keys_rows = [
    {"wine_id": 1},
    {"wine_id": 3}
]
result = client.read_table(name=table_name,database_name=database_name, is_online=True, online_config=online_config, entity_row=primary_keys_rows)
result.show()

# Read offline feature table--DLC on iceberg(enable catalog)
# read offline feature table--DLC on iceberg(enable catalog)
get_df = client.read_table(name=table_name, database_name=database_name, catalog_name=catalog_name)
get_df.show(30)

# Get online feature table data--DLC on iceberg(enable catalog)
# read online feature table--DLC on iceberg(enable catalog)
primary_keys_rows = [
    {"wine_id": 1},
    {"wine_id": 3}
]
result = client.read_table(name=table_name,database_name=database_name, is_online=True, online_config=online_config, entity_row=primary_keys_rows, catalog_name=catalog_name)
result.show()

# Get feature table data--EMR on hive--&--DLC on iceberg(catalog not enabled)
# get feature table data--EMR on hive--&--DLC on iceberg(disable catalog)
featureTable = client.get_table(name=table_name, database_name=database_name)

# Get online feature table data--DLC on iceberg(enable catalog)
# get feature table data--DLC on iceberg(enable catalog)
featureTable = client.get_table(name=table_name, database_name=database_name, catalog_name=catalog_name)

# drop feature table data--EMR on hive--&--DLC on iceberg(disable catalog)
# drop feature table data--EMR on hive--&--DLC on iceberg(disable catalog)
client.drop_table(table_name)

# Delete feature table--DLC on iceberg(enable catalog)
# drop feature table data--DLC on iceberg(enable catalog)
client.drop_table(table_name,catalog_name="DataLakeCatalog")

from wedata.feature_store.constants.constants import APPEND

# Batch processing write--EMR on hive--&--DLC on iceberg(catalog not enabled)
# batch write--EMR on hive--&--DLC on iceberg(disable catalog)
client.write_table(
    name="user_features",
    df=user_data_df,
    database_name="feature_db",
    mode="append"
)

# Stream processing write--EMR on hive--&--DLC on iceberg(catalog not enabled)
# strem write--EMR on hive--&--DLC on iceberg(disable catalog)
streaming_query = client.write_table(
    name="user_features",
    df=streaming_df,
    database_name="feature_db",
    checkpoint_location="/checkpoints/user_features",
    trigger={"processingTime": "10 seconds"}
)

# Batch processing write--DLC on iceberg(catalog enabled)
# batch write--EMR on hive--&--DLC on iceberg(enable catalog)
client.write_table(
    name="user_features",
    df=user_data_df,
    database_name="feature_db",
    mode="append",
    catalog_name=catalog_name
)

# Stream processing write--DLC on iceberg(catalog enabled)
# strem write--EMR on hive--&--DLC on iceberg(enable catalog)
streaming_query = client.write_table(
    name="user_features",
    df=streaming_df,
    database_name="feature_db",
    checkpoint_location="/checkpoints/user_features",
    trigger={"processingTime": "10 seconds"},
    catalog_name=catalog_name
)

from wedata.feature_store.entities.feature_lookup import FeatureLookup

# Feature search--EMR on hive--&--DLC on iceberg(catalog not enabled)--&--DLC on iceberg(catalog enabled)
# feature lookup--EMR on hive--&--DLC on iceberg(disable catalog)--&--DLC on iceberg(enable catalog)
wine_feature_lookup = FeatureLookup(
    table_name=table_name,
    lookup_key="wine_id",
    timestamp_lookup_key="event_timestamp"
)

from wedata.feature_store.entities.feature_lookup import FeatureLookup
from wedata.feature_store.entities.feature_function import FeatureFunction
# Define feature search
# define feature lookup
wine_feature_lookup = FeatureLookup(
    table_name=table_name,
    lookup_key="wine_id",
    timestamp_lookup_key="event_timestamp"
)

# Build training data
# prepare training data
inference_data_df = wine_df.select(f"wine_id", "quality", "event_timestamp")

# create trainingset--EMR on hive--&--DLC on iceberg(disable catalog)
# create trainingset--EMR on hive--&--DLC on iceberg(disable catalog)
training_set = client.create_training_set(
    df=inference_data_df,  # basic data
    feature_lookups=[wine_feature_lookup],  # feature search configuration
    label="quality",  # tag column
    exclude_columns=["wine_id", "event_timestamp"]  # exclude unnecessary columns
)

# create trainingset--DLC on iceberg(enable catalog)
# create trainingset--DLC on iceberg(enable catalog)
training_set = client.create_training_set(
    df=inference_data_df,  # basic data
    feature_lookups=[wine_feature_lookup],  # feature search configuration
    label="quality",  # tag column
    exclude_columns=["wine_id", "event_timestamp"],  # exclude unnecessary columns
    database_name=database_name,
    catalog_name=catalog_name
)

# get final training DataFrame
# get final training DataFrame
training_df = training_set.load_df()

# Print training set data
# print trainingset data
print(f"\\n=== Training set data ===")
training_df.show(10, True)

# Training model
# train model
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
from sklearn.ensemble import RandomForestClassifier
import mlflow.sklearn
import pandas as pd
import os 
project_id=os.environ["WEDATA_PROJECT_ID"]
mlflow.set_experiment(experiment_name=expirement_name)

# Set mlflow information reporting address
# set mlflow tracking_uri
mlflow.set_tracking_uri("http://30.22.36.75:5000")

# Convert Spark DataFrame to Pandas DataFrame for training
# convert Spark DataFrame to Pandas DataFrame to train
train_pd = training_df.toPandas()

# Delete timestamp column
# delete timestamp
# train_pd.drop('event_timestamp', axis=1)

# Prepare features and tags
# prepare features and tags
X = train_pd.drop('quality', axis=1)
y = train_pd['quality']

# Divide the training set and test set
# split traningset and testset
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# Convert date to timestamp (seconds)
# or convert datatime to timestamp(second)
for col in X_train.select_dtypes(include=['datetime', 'datetimetz']):
    X_train[col] = X_train[col].astype('int64') // 10**9   # Convert nanoseconds to seconds

# Confirm no missing values cause dtype to downgrade to object
# Verify that no missing values cause the dtype to be downgraded to object.
X_train = X_train.fillna(X_train.median(numeric_only=True))

# initialize and train model and record model--EMR on hive--&--DLC on iceberg(disable catalog)--&--DLC on iceberg(enable catalog)
# Initialize, train and log the model.--EMR on hive--&--DLC on iceberg(disable catalog)--&--DLC on iceberg(enable catalog)
model = RandomForestClassifier(n_estimators=100, max_depth=3, random_state=42)
model.fit(X_train, y_train)

with mlflow.start_run():
    client.log_model(
        model=model,
        artifact_path="wine_quality_prediction",  # model artifact path
        flavor=mlflow.sklearn,
        training_set=training_set,
        registered_model_name=model_name,  # model name (if catalog is enabled, must be catalog model name)
    )

# run score_batch--EMR on hive--&--DLC on iceberg(disable catalog)--&--DLC on iceberg(enable catalog)
# run score_batch--EMR on hive--&--DLC on iceberg(disable catalog)--&--DLC on iceberg(enable catalog)
result = client.score_batch(model_uri=f"models:/{model_name}/1", df=wine_df)
result.show()