# -*- coding: utf-8 -*-
import dataiku

from clv_forecast.utils.clustering import apply_clustering

# Read recipe inputs
inference_data_prepared = dataiku.Dataset("inference_data_prepared")
inference_data_prepared_df = inference_data_prepared.get_dataframe()


inference_data_prepared_df["future_clv_cluster_predicted"] = apply_clustering(
    inference_data_prepared_df["prediction_ml_future_monthly_average_value"]
)
inference_data_prepared_df["future_clv_cluster_lifetime"] = apply_clustering(
    inference_data_prepared_df["prediction_lifetime_future_monthly_average_value"]
)
inference_data_prepared_df["low_transaction_number_customer"] = (
    inference_data_prepared_df["low_transaction_number_customer"]
    .astype(str)
    .str.lower()
)

inference_data_clustered_df = inference_data_prepared_df


# Write recipe outputs
inference_data_clustered = dataiku.Dataset("inference_data_clustered")
inference_data_clustered.write_dataframe(
    inference_data_clustered_df, infer_schema=False
)
