# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
# -*- coding: utf-8 -*-
import dataiku
import pandas as pd, numpy as np

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
from dku_utils.projects.recipes.recipe_commons import adapt_recipe_engine_to_priority_and_availability, update_recipe_ouput_schema
from dku_utils.projects.project_commons import get_current_project_and_variables, get_project_and_variables

from dku_utils.projects.folders.pickles.folder_pickles import write_pickle_in_managed_folder
from markdown_optimization.config.functions import compute_all_hierarchy_dependencies

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
project, variables = get_current_project_and_variables()
global_variables = variables["standard"]
forecast_time_granularity = "week"

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
granularity_x_hierarchy_constraints_df = dataiku.Dataset("granularity_x_hierarchy_constraints").get_dataframe()
dates_boundaries_df = dataiku.Dataset("dates_boundaries").get_dataframe()
demand_forecast_date_df = dataiku.Dataset("demand_forecast_dates").get_dataframe()

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
products_hierarchy = ["product_id"]
for column_name in products_hierarchy:
    granularity_x_hierarchy_constraints_df[column_name] = granularity_x_hierarchy_constraints_df[column_name].astype(str)

granularity_x_hierarchy_constraints_df.rename({"forecast_granularity": "granularity"}, axis=1, inplace=True)

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
all_granularities = list(np.unique(granularity_x_hierarchy_constraints_df["granularity"]))
all_products = list(np.unique(granularity_x_hierarchy_constraints_df["product_id"]))

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
forecast_horizon = 10
granularity_x_hierarchy_constraints = {}
for granularity in all_granularities:
    print("Computing granularity '{}' x product hierarchy constraints".format(granularity))
    granularity_x_hierarchy_constraints[granularity] = {}
    granularity_df = granularity_x_hierarchy_constraints_df[granularity_x_hierarchy_constraints_df["granularity"]==granularity].copy()
    for column in products_hierarchy:
        granularity_hiearchical_layer_df = granularity_df[["granularity", column]].drop_duplicates()
        granularity_x_hierarchy_constraints[granularity][column] = list(np.unique(granularity_hiearchical_layer_df[column]))

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
write_pickle_in_managed_folder(project, "webapp_folder", granularity_x_hierarchy_constraints, "granularity_x_hierarchy_constraints")

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
hierarchy_dependencies =\
compute_all_hierarchy_dependencies(granularity_x_hierarchy_constraints_df[products_hierarchy], products_hierarchy, "product_id")
write_pickle_in_managed_folder(project, "webapp_folder", hierarchy_dependencies, "hierarchy_dependencies")
# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
last_available_date = demand_forecast_date_df["date"][0].strftime('%Y-%m-%d')
demand_forecast_date = demand_forecast_date_df["forecast_date"][0].strftime('%Y-%m-%d')
webapp_min_date = dates_boundaries_df["min_date"][0].strftime('%Y-%m-%d')
max_date = dates_boundaries_df["max_date"][0]

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
webapp_max_date = max_date.strftime('%Y-%m-%d')

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
dates = list(pd.date_range(start=webapp_min_date, end=webapp_max_date, freq="D"))
dates = [date.strftime('%Y-%m-%d') for date in dates]

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
date_constraints = {
        "min_date": webapp_min_date,
        "max_date": webapp_max_date,
        "demand_forecast_date": demand_forecast_date,
        "current_date": last_available_date,
        "dates": dates
    }

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
write_pickle_in_managed_folder(project, "webapp_folder", date_constraints, "date_constraints")