import dataiku

from clv_forecast.constants import (
    CATEGORICAL_VALUES_PROCESSING,
    NUMERICAL_VALUES_PROCESSING,
)
from dku_utils.projects.visual_ml.visual_ml_commons import (
    remove_unavailable_features_from_ml_task_features_handling,
    update_ml_task_feature_preprocessing,
    enable_features_in_ml_task_settings,
)


def update_processing_features(
    project, variables, ml_task, ml_task_settings, train_dataset_name
):
    remove_unavailable_features_from_ml_task_features_handling(
        project, ml_task_settings, train_dataset_name
    )
    enable_features = []
    preprocessing_dict = dict()

    if variables["standard"]["leverage_rfm_segmentation_app"]:
        rfm_dataset = dataiku.Dataset(variables["standard"]["rfm_dataset_app"])
        schema = rfm_dataset.read_schema()

        for column in schema:
            if (column["name"] != "customer_id") and (
                column["name"] != "rfm_reference_date"
            ):
                if column["type"] == "string":
                    preprocessing_dict.update(
                        {column["name"]: CATEGORICAL_VALUES_PROCESSING}
                    )
                else:
                    preprocessing_dict.update(
                        {column["name"]: NUMERICAL_VALUES_PROCESSING}
                    )
                enable_features.append(column["name"])

    if variables["standard"]["leverage_customer_metadata_app"]:
        rfm_dataset = dataiku.Dataset(
            variables["standard"]["customer_metadata_dataset_app"]
        )

        schema = rfm_dataset.read_schema()
        columns = list(
            filter(
                lambda name: name != "customer_id",
                map(lambda pair: pair["name"], schema),
            )
        )
        for column in columns:
            if column in variables["standard"]["categorical_metadata_app"]:
                preprocessing_dict.update({column: CATEGORICAL_VALUES_PROCESSING})
                enable_features.append(column)
            elif column in variables["standard"]["numerical_metadata_app"]:
                preprocessing_dict.update({column: NUMERICAL_VALUES_PROCESSING})
                enable_features.append(column)
            else:
                pass

    for column, preprocessing in preprocessing_dict.items():
        update_ml_task_feature_preprocessing(ml_task_settings, column, preprocessing)

    enable_features_in_ml_task_settings(
        ml_task_settings, list_of_features_to_enable=enable_features
    )
