# Python script to compute advance working dates

# Weekend Example: using Saturday or Sunday and add 2 working days:
# Monday = 1 working day, Tuesday = 1 working day, Wednesday = Delivery Date
# Assumption for weekends, base date is NOT considered as one working day

# Weekday Example: using Monday and adding 2 working days:
# Monday = 1 working day, Tuesday = 1 working day, Wednesday = Delivery Date
# Assumption for weekdays, base date is considered as one working day

# Forecasting on future dates without information on holidays will not be accurate.


import pandas as pd
from datetime import datetime as dt, timedelta


# Function to check if a given date is a weekend [Saturday(5) or Sunday(6)]
def is_weekend(date):
    return date.weekday() in (5, 6)


# Function to compute the in advance working days
def compute_working_days(start_date, num_working_days, holiday_list, after_cut_off = 0):
    # after_cut_off = 0 for orders before 4pm
    # after_cut_off = 1 for orders after 4pm

    counter = 0

    # For dates on weekend or holiday or after cut off time, additional 1 working day to the loop
    if (after_cut_off == 1) or is_weekend(start_date) or (start_date in holiday_list):
        counter -= 1

    # Loop to increase n number of working days, if weekend/holiday, skip counter
    while counter < num_working_days:
        start_date = start_date + timedelta(days=1)

        if is_weekend(start_date) or (start_date in holiday_list):
            continue

        counter += 1

    return start_date


# Main function
def model(dbt, session):

    # Setting configuration
    dbt.config(materialized="table",
                packages = ["pandas"])

    # Import data from upstream dbt model
    sp_df_date = dbt.ref("int__dates")

    # Filter holiday dates into a dataframe
    sp_df_filter = sp_df_date.filter(sp_df_date['IS_COMPANY_HOLIDAY'] == 1)
    sp_df_holiday_date = sp_df_filter.select('DATE_DAY')

    # Convert snowpark dataframe to pandas dataframe
    pd_df_date = sp_df_date.to_pandas()
    pd_df_holiday_date = sp_df_holiday_date.to_pandas()
    
    # Store holiday dates in a list
    holiday_list = pd_df_holiday_date['DATE_DAY'].tolist()

    # Applying add working days function to df
    add_working_days = [1, 2, 3, 7, 30, 90 ,365]

    for working_day in add_working_days:
        for after_cut_off in (0, 1):

            if after_cut_off == 0:
                col_name = f'ADD_{working_day}_WORKING_DAY_INCLUDED_BASE_DATE'
            else:
                col_name = f'ADD_{working_day}_WORKING_DAY_EXCLUDED_BASE_DATE'

            # Apply function and append new calculated columns in the dataframe
            pd_df_date[col_name] = pd_df_date['DATE_DAY'].apply(lambda start_date: compute_working_days(start_date, working_day, holiday_list, after_cut_off))

    return pd_df_date