import dataiku

from dku_utils.projects.project_commons import get_current_project_and_variables
from dku_utils.projects.recipes.join_recipe import programmaticJoinHandler
from dku_utils.projects.recipes.recipe_commons import (
    adapt_recipe_engine_to_priority_and_availability,
    update_recipe_ouput_schema,
)


def join_uplift():
    project, variables = get_current_project_and_variables()
    recipe_name = "compute_predicted_demand_with_reference"

    base_dataset = dataiku.Dataset("predicted_demand_prepared")

    schema = base_dataset.read_schema()
    main_columns = list(map(lambda pair: pair["name"], schema))

    reference_columns = [
        "profit_per_unit",
        "margin",
        "forecasted_volume",
        "revenue",
        "sellthrought_rate",
        "total_profit",
        "lost_sales",
        "lost_profit",
    ]
    constant_strategy_alias = {
        "profit_per_unit": "profit_per_unit_static",
        "margin": "margin_static",
        "forecasted_volume": "forecasted_volume_static",
        "revenue": "revenue_static",
        "sellthrought_rate": "sellthrought_rate_static",
        "total_profit": "total_profit_static",
        "lost_sales": "lost_sales_static",
        "lost_profit": "lost_profit_static",
    }
    zero_strategy_alias = {
        "profit_per_unit": "profit_per_unit_zero",
        "margin": "margin_zero",
        "forecasted_volume": "forecasted_volume_zero",
        "revenue": "revenue_zero",
        "sellthrought_rate": "sellthrought_rate_zero",
        "total_profit": "total_profit_zero",
        "lost_sales": "lost_sales_zero",
        "lost_profit": "lost_profit_zero",
    }

    join_handler = programmaticJoinHandler(
        project=project,
        recipe_name=recipe_name,
        main_dataset_name="predicted_demand_prepared",
        main_dataset_columns_to_select=main_columns,
        main_dataset_columns_to_select_alias={},
        main_dataset_computed_columns=None,
    )

    join_handler.add_one_join_on_main_dataset(
        dataset_to_join_name="sales_x_features_reference_static_scored_prepared",
        dataset_to_join_columns_to_select=reference_columns,
        join_type="LEFT",
        columns_prefix=None,
        left_join_key=["product_id", "forecast_granularity"],
        right_join_key=["product_id", "forecast_granularity"],
        columns_to_select_alias=constant_strategy_alias,
        dataset_computed_columns=None,
    )

    join_handler.add_one_join_on_main_dataset(
        dataset_to_join_name="sales_x_features_zero_scored_prepared",
        dataset_to_join_columns_to_select=reference_columns,
        join_type="LEFT",
        columns_prefix=None,
        left_join_key=["product_id", "forecast_granularity"],
        right_join_key=["product_id", "forecast_granularity"],
        columns_to_select_alias=zero_strategy_alias,
        dataset_computed_columns=None,
    )
    adapt_recipe_engine_to_priority_and_availability(project, recipe_name)
    update_recipe_ouput_schema(project, recipe_name)
