%pip install tencent-wedata-feature-engineering

from wedata.feature_store.client import FeatureStoreClient
# 构建特征工程客户端实例
# 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
# 创建用户特征表--EMR on hive
# create user's feature table--EMR on hive
feature_table = client.create_table(
    name=table_name,  # 表名
    primary_keys=["wine_id"],      # 主键
    df=wine_features_df,           # 数据
    engine_type=EngineTypes.HIVE_ENGINE,
    data_source_name=data_source_name,
    timestamp_key="event_timestamp",
    tags={ # 业务标签
        "purpose": "demo", 
        "create_by": "wedata"
    }
)

# 创建用户特征表--DLC on iceberg（不开启catalog）
# create user's feature table--DLC on iceberg（disable catalog）
feature_table = client.create_table(
    name=table_name,  # 表名
    database_name=database_name,
    primary_keys=["wine_id"],      # 主键
    df=wine_features_df,           # 数据
    engine_type=EngineTypes.ICEBERG_ENGINE,
    data_source_name=data_source_name,
    timestamp_key="event_timestamp",
    tags={ # 业务标签
        "purpose": "demo", 
        "create_by": "wedata"
    }
)

# 创建用户特征表--DLC on iceberg（开启catalog）
# create user's feature table--DLC on iceberg（enable catalog）
feature_table = client.create_table(
    name=table_name,  # 表名
    primary_keys=["wine_id"],      # 主键
    df=wine_features_df,           # 数据
    engine_type=EngineTypes.ICEBERG_ENGINE,
    data_source_name=data_source_name,
    timestamp_key="event_timestamp",
    tags={ # 业务标签
        "purpose": "demo", 
        "create_by": "wedata"
    },
    catalog_name=catalog_name
)

from wedata.feature_store.constants.engine_types import EngineTypes
# 注册特征表--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",])

# 注册特征表--DLC on iceberg（不开启catalog）
# 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",])

# 注册特征表--DLC on iceberg（开启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)

# 读取离线特征表--EMR on hive--&--DLC on iceberg(不开启catalog)
# 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)

# 读取在线特征表--EMR on hive--&--DLC on iceberg(不开启catalog)
# 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()

# 读取离线特征表--DLC on iceberg(开启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)

# 读取在线特征表--DLC on iceberg(开启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()

# 获取特征表数据--EMR on hive--&--DLC on iceberg(不开启catalog)
# get feature table data--EMR on hive--&--DLC on iceberg(disable catalog)
featureTable = client.get_table(name=table_name, database_name=database_name)

# 读取在线特征表--DLC on iceberg(开启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)

# 删除特征表--EMR on hive--&--DLC on iceberg(不开启catalog)
# drop feature table data--EMR on hive--&--DLC on iceberg(disable catalog)
client.drop_table(table_name)

# 删除特征表--DLC on iceberg(开启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

# 批处理写入--EMR on hive--&--DLC on iceberg(不开启catalog)
# 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"
)

# 流式处理写入--EMR on hive--&--DLC on iceberg(不开启catalog)
# 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"}
)

# 批处理写入--DLC on iceberg(开启catalog)
# 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
)

# 流式处理写入--DLC on iceberg(开启catalog)
# 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

# 特征查找--EMR on hive--&--DLC on iceberg(不开启catalog)--&--DLC on iceberg(开启catalog)
# 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 lookup
wine_feature_lookup = FeatureLookup(
    table_name=table_name,
    lookup_key="wine_id",
    timestamp_lookup_key="event_timestamp"
)

# 构建训练数据
# prepare training data
inference_data_df = wine_df.select(f"wine_id", "quality", "event_timestamp")

# 创建训练集--EMR on hive--&--DLC on iceberg(不开启catalog)
# create trainingset--EMR on hive--&--DLC on iceberg(disable catalog)
training_set = client.create_training_set(
    df=inference_data_df,  # 基础数据
    feature_lookups=[wine_feature_lookup],  # 特征查找配置
    label="quality",  # 标签列
    exclude_columns=["wine_id", "event_timestamp"]  # 排除不需要的列
)

# 创建训练集--DLC on iceberg(开启catalog)
# create trainingset--DLC on iceberg(enable catalog)
training_set = client.create_training_set(
    df=inference_data_df,  # 基础数据
    feature_lookups=[wine_feature_lookup],  # 特征查找配置
    label="quality",  # 标签列
    exclude_columns=["wine_id", "event_timestamp"],  # 排除不需要的列
    database_name=database_name,
    catalog_name=catalog_name
)

# 获取最终的训练DataFrame
# get final training DataFrame
training_df = training_set.load_df()

# 打印训练集数据
# print trainingset data
print(f"\\n=== 训练集数据 ===")
training_df.show(10, True)

# 训练模型
# 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)

# 设置mlflow信息上报地址
# set mlflow tracking_uri
mlflow.set_tracking_uri("http://30.22.36.75:5000")

# 将Spark DataFrame转换为Pandas DataFrame用于训练
# convert Spark DataFrame to Pandas DataFrame to train
train_pd = training_df.toPandas()

# 删除时间戳列
# delete timestamp
# train_pd.drop('event_timestamp', axis=1)

# 准备特征和标签
# prepare features and tags
X = train_pd.drop('quality', axis=1)
y = train_pd['quality']

# 划分训练集和测试集
# split traningset and testset
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# 或者把日期转成时间戳（秒）
# 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   # 纳秒→秒

# 确认没有缺失值导致 dtype 被降级为 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))

# 初始化并训练模型和记录模型--EMR on hive--&--DLC on iceberg(不开启catalog)--&--DLC on iceberg(开启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,  # 模型名称（开启catalog后需要写catalog的模型名称）model name(if enable catalog, must be catalog model name)
    )

# 运行score_batch--EMR on hive--&--DLC on iceberg(不开启catalog)--&--DLC on iceberg(开启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()

from wedata.feature_store.common.store_config.redis import RedisStoreConfig
from wedata.feature_store.cloud_sdk_client.models import TaskSchedulerConfiguration

online_config = RedisStoreConfig(host=redis_host, port=redis_password, db=redis_db, password=redis_password)

# 将离线特征表发布成为在线特征表--EMR on hive--&--DLC on iceberg(不开启catalog)
# publish a offline feature table to online feature table--EMR on hive--&--DLC on iceberg(disable catalog)

# Case 1 将离线表table_name通过指定离线数据源传入到默认在线数据源，并采用快照方式
# Case 1 Pass the offline table table_name from the specified offline data source to the default online data source in snapshot mode.
result = client.publish_table(table_name=table_name, data_source_name=data_source_name, database_name=database_name,
                              is_cycle=False, is_use_default_online=True)

# Case 2 将离线表table_name通过指定离线数据源传入到指定数据源(非默认)，并采用快照方式, 需注意这里的IP要与录入的数据源中登录的数据源IP信息一致。
# Case 2 Sync the offline table table_name from the specified offline data source to the designated non-default data source in snapshot mode. Note that the IP address here must be consistent with the login IP information of the entered data source.
result = client.publish_table(name=table_name, data_source_name=data_source_name, database_name=database_name,
                               is_cycle=False, is_use_default_online=False, online_config=online_config)

# Case 3 将离线表table_name通过指定离线数据源传入到默认在线数据源，并采用周期的方式,每五分钟运行一次
# Sync the offline table table_name from the specified offline data source to the default online data source in cycle mode, with a run interval of every five minutes.
cycle_obj = TaskSchedulerConfiguration()
cycle_obj.CrontabExpression = "0 0/5 * * * ?"
result = client.publish_table(name=table_name, data_source_name=data_source_name, database_name=database_name,
                              is_cycle=True, cycle_obj=cycle_obj, is_use_default_online=True)
                           

# 将离线特征表发布成为在线特征表--DLC on iceberg(开启catalog)
# publish a offline feature table to online feature table--DLC on iceberg(enable catalog)

# Case 1 将离线表table_name通过指定离线数据源传入到默认在线数据源，并采用快照方式
# Case 1 Pass the offline table table_name from the specified offline data source to the default online data source in snapshot mode.
result = client.publish_table(table_name=table_name, data_source_name=data_source_name, database_name=database_name,
                              is_cycle=False, is_use_default_online=True, catalog_name=catalog_name)

# Case 2 将离线表table_name通过指定离线数据源传入到指定数据源(非默认)，并采用快照方式, 需注意这里的IP要与录入的数据源中登录的数据源IP信息一致。
# Case 2 Sync the offline table table_name from the specified offline data source to the designated non-default data source in snapshot mode. Note that the IP address here must be consistent with the login IP information of the entered data source.
result = client.publish_table(table_name=table_name, data_source_name=data_source_name, database_name=database_name,
                              is_cycle=False, is_use_default_online=False, online_config=online_config, catalog_name=catalog_name)

# Case 3 将离线表table_name通过指定离线数据源传入到默认在线数据源，并采用周期的方式,每五分钟运行一次
# Sync the offline table table_name from the specified offline data source to the default online data source in cycle mode, with a run interval of every five minutes.
cycle_obj = TaskSchedulerConfiguration()
cycle_obj.CrontabExpression = "0 0/5 * * * ?"
result = client.publish_table(table_name=table_name, data_source_name=data_source_name, database_name=database_name,
                              is_cycle=True, cycle_obj=cycle_obj, is_use_default_online=True, catalog_name=catalog_name)