%pip install tencent-wedata-feature-engineering

%pip list | grep 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)
    )