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

from clv_forecast.utils.clustering import apply_clustering

# Read recipe inputs
ml_data_test_scored = dataiku.Dataset("ml_data_test_scored_prepared")
ml_data_test_scored_df = ml_data_test_scored.get_dataframe()


ml_data_test_scored_df["future_clv_cluster_predicted"] = apply_clustering(
    ml_data_test_scored_df["prediction_ml_future_monthly_average_value"]
)
ml_data_test_scored_df["future_clv_cluster_lifetime"] = apply_clustering(
    ml_data_test_scored_df["prediction_lifetime_future_monthly_average_value"]
)

ml_data_test_clustered_df = (
    ml_data_test_scored_df  # For this sample code, simply copy input to output
)


# Write recipe outputs
ml_data_test_clustered = dataiku.Dataset("ml_data_test_clustered")
ml_data_test_clustered.write_with_schema(ml_data_test_clustered_df)
