# Check whether the Feature Engineering SDK is installed (it is preinstalled in the Machine Learning image and the version will be printed).
%pip list | grep wedata3-feature-engineering

# If not installed, uncomment the next line and run it (requires kernel Python 3.10 or later and MLflow 3.0.0 or later). Restart the kernel after installation.
# %pip install "wedata3-feature-engineering[mlflow3]==0.2.3"

import pandas as pd
from pyspark.sql.functions import current_timestamp
from wedata.feature_engineering.client import FeatureEngineeringClient
from wedata.feature_store.constants.engine_types import EngineTypes

# The Studio (DLC) Notebook kernel has a built-in spark session, so you can use the spark variable directly without manually creating a SparkSession.

# Initialize the Feature Engineering client (authentication information is automatically injected from environment variables by the Notebook kernel, so no manual configuration is required).
fe = FeatureEngineeringClient(spark)

# Read the uploaded CSV and add a timestamp column.
# Note: Replace the path below with the actual directory path where you uploaded the CSV in step 1.1.
# To obtain the full path, right-click the CSV file in the file manager on the left and select Copy Path.
# Example: /Workspace/Users/<your account>/customer_churn_class.csv
pdf = pd.read_csv("/Workspace/Users/<your account>/customer_churn_class.csv")
df = spark.createDataFrame(pdf).withColumn("event_timestamp", current_timestamp())

# Create a feature table: specify the primary key and timestamp key, set engine_type to Iceberg, and use df to infer the schema and write the initial data.
fe.create_table(
    name="customer_churn_features",          # (Modifiable) Feature table name
    catalog_name="your_catalog",          # Target Catalog name
    database_name="your_schema",           # Target Schema name
    primary_keys=["customer_id"],            # Primary key (composite keys supported)
    timestamp_key="event_timestamp",         # Timestamp key (singular)
    engine_type=EngineTypes.ICEBERG_ENGINE,
    df=df,                                    # Use df to infer the schema and write the initial data.
    description="Customer churn feature table for AutoML training.",
)

# Read back the offline feature table to verify the write result. The table name is the three-part full name from create_table above. Replace it with your own catalog.schema.
fe.read_table(name="your_catalog.your_schema.customer_churn_features").show(5)

curl -X POST "https://<service-url>/predict" \\
  -H "Authorization: Bearer <token>" \\
  -H "Content-Type: application/json" \\
  -d '{"dataframe_records": [{
            "age": 27,
            "gender": "Male",
            "tenure_months": 70,
            "monthly_charges": 112.99,
            "total_charges": 7410.38,
            "contract_type": "One Year",
            "payment_method": "Credit Card",
            "num_support_tickets": 5,
            "avg_monthly_usage": 271.91
        }]}'

%pip install "wedata3-mlflow-tcdeploy-plugin"

from mlflow.deployments import get_deploy_client

# target is fixed to "tcdeploy".
client = get_deploy_client("tcdeploy")

# For services deployed from the console, you can obtain the service group ID by calling list_deployments(). The name field in the returned item is the service_group_id.
service_group_id = client.list_deployments()[0]["name"]

# For the deployment_name parameter of predict, pass the service group ID (service_group_id, in the format msg-xxxxxx), not the service name.
result = client.predict(
    deployment_name=service_group_id,
    inputs={"dataframe_records": [{
            "age": 27,
            "gender": "Male",
            "tenure_months": 70,
            "monthly_charges": 112.99,
            "total_charges": 7410.38,
            "contract_type": "One Year",
            "payment_method": "Credit Card",
            "num_support_tickets": 5,
            "avg_monthly_usage": 271.91
        }]},
   config={"auth_token": "Bearer <token>"},   # Required when authentication is enabled for the service.
)
print(result)