try:

    from datetime import timedelta
    from airflow import DAG
    
    from airflow.operators.python_operator import PythonOperator
    from airflow.operators.http_operator import SimpleHttpOperator
    from datetime import datetime
    from pandas.io.json import json_normalize
    from airflow.operators.postgres_operator import PostgresOperator

    import pandas as pd
    import json
    import requests
    import numpy as np
    import re
    import calendar

    import psycopg2
    from sqlalchemy import create_engine, inspect, text

except Exception as e:
    print("Error {} ".format(e))

dRoW_api_end_url = "https://drow.cloud"

def getDrowToken(**context):

    response = requests.post(
        url=f"{dRoW_api_end_url}/api/auth/authenticate",
        data={
            "username": "icwp2@drow.cloud",
            "password": "dGVzdDAxQHRlc3QuY29t"
        }
    ).json()
    context["ti"].xcom_push(key="token", value=response['token'])
    # return 'DLLM{}'.format(response)


def getMongoDB(**context):
    token = context.get("ti").xcom_pull(key="token")
    host                  = 'drowdatewarehouse.crlwwhgepgi7.ap-east-1.rds.amazonaws.com'  
    # User name of the database server
    dbUserName            = 'dRowAdmin'  
    # Password for the database user
    dbUserPassword        = 'drowsuper'  
    # Name of the database 
    database              = 'drowDateWareHouse'
    # Character set
    charSet               = "utf8mb4"  
    port                  = "5432"
    conn_string = ('postgres://' +
                           dbUserName + ':' + 
                           dbUserPassword +
                           '@' + host + ':' + port +
                           '/' + database)
    
    db = create_engine(conn_string)
    conn = db.connect()

    mappings_risk_reg = {
        "Date_of_Closure_of_Early_Warning": "Date_of_Close_of_EW",
        "Date_Notified": "Date_of_Early_Warning",
        "Notified_by___PM_or_C": "Notified_by",
        "Status___Live__Closed": "Status",
    }
    mappings_nec = {
        "Incident_No": "NEC_Event_No",
        "CE_No": "CE_No",
        "CE_Increase___Decrease": "CE_Increase_Decrease",
    }

    def update_months(sourcedate, months):
        month = sourcedate.month - 1 + months
        year = sourcedate.year + month // 12
        month = month % 12 + 1
        day = min(sourcedate.day, calendar.monthrange(year,month)[1])
        return datetime(year, month, day)

    with conn as conn:
        # Get current date
        today = datetime.now()

        # Load data from SQL into DataFrame
        # Load c1 - p5 data
        df_c1 = pd.read_sql('SELECT * FROM public."_6wsd21_edms_nec_cr_c1" WHERE "Doc_Date" IS NOT NULL;', conn)
        df_c1['Table'] = 'c1'
        df_c1.rename(columns=mappings_nec, inplace=True)
        df_c1['Doc_Date'] = df_c1['Doc_Date'].astype(str)

        df_c2 = pd.read_sql('SELECT * FROM public."_6wsd21_edms_nec_cr_c2" WHERE "Doc_Date" IS NOT NULL;', conn)
        df_c2['Table'] = 'c2'
        df_c2.rename(columns=mappings_nec, inplace=True)
        df_c2['Doc_Date'] = df_c2['Doc_Date'].astype(str)

        df_c3 = pd.read_sql('SELECT * FROM public."_6wsd21_edms_nec_cr_c3" WHERE "Doc_Date" IS NOT NULL;', conn)
        df_c3['Table'] = 'c3'
        df_c3.rename(columns=mappings_nec, inplace=True)
        df_c3['Doc_Date'] = df_c3['Doc_Date'].astype(str)

        df_c4 = pd.read_sql('SELECT * FROM public."_6wsd21_edms_nec_cr_c4" WHERE "Doc_Date" IS NOT NULL;', conn)
        df_c4['Table'] = 'c4'
        df_c4.rename(columns=mappings_nec, inplace=True)
        df_c4['Doc_Date'] = df_c4['Doc_Date'].astype(str)

        df_p1 = pd.read_sql('SELECT * FROM public."_6wsd21_edms_nec_cr_p1" WHERE "Doc_Date" IS NOT NULL;', conn)
        df_p1['Table'] = 'p1'
        df_p1.rename(columns=mappings_nec, inplace=True)
        df_p1['Doc_Date'] = df_p1['Doc_Date'].astype(str)

        df_p2 = pd.read_sql('SELECT * FROM public."_6wsd21_edms_nec_cr_p2" WHERE "Doc_Date" IS NOT NULL;', conn)
        df_p2['Table'] = 'p2'
        df_p2.rename(columns=mappings_nec, inplace=True)
        df_p2['Doc_Date'] = df_p2['Doc_Date'].astype(str)
        
        # Note: uses table _6wsd21 not 6wsd21, need to be updated in future
        # df_p3 = pd.read_sql('SELECT * FROM public."_6wsd21_edms_nec_cr_p3" WHERE "Doc_Date" IS NOT NULL;', conn)
        df_p3 = pd.DataFrame()
        if not df_p3.empty:
            df_p3['Table'] = 'p3'
            df_p3.rename(columns=mappings_nec, inplace=True)
            df_p3 = df_p3[df_p3['CE_No'] != 'null0']
            df_p3['Doc_Date'] = df_p3['Doc_Date'].astype(str)

        df_p4 = pd.read_sql('SELECT * FROM public."_6wsd21_edms_nec_cr_p4" WHERE "Doc_Date" IS NOT NULL;', conn)
        if not df_p4.empty:
            df_p4['Table'] = 'p4'
            df_p4.rename(columns=mappings_nec, inplace=True)
            df_p4['Doc_Date'] = df_p4['Doc_Date'].astype(str)

        df_p5 = pd.read_sql('SELECT * FROM public."_6wsd21_edms_nec_cr_p5" WHERE "Doc_Date" IS NOT NULL;', conn)
        if not df_p5.empty:
            df_p5['Table'] = 'p5'
            df_p5.rename(columns=mappings_nec, inplace=True)
            df_p5['Doc_Date'] = df_p5['Doc_Date'].astype(str)

        df_all = pd.concat([df_c1, df_c2, df_c3, df_c4, df_p1, df_p2, df_p3], axis=0, ignore_index=True, sort=False)

        df_events = pd.read_sql('SELECT * FROM public."_6wsd21_edms_nec_cr_event";', conn)
        # No table for PCD, uncomment when data is in
        # df_from_sql_nec_section_of_works = pd.read_sql('SELECT * FROM public."_6wsd21_edms_nec_section_of_work";', conn)
        # df_from_sql_PCD = pd.read_sql('SELECT * FROM public."_6wsd21_nec_key_date_data";', conn)
        df_from_sql_nec_section_of_works = pd.DataFrame()
        df_from_sql_PCD = pd.DataFrame()

        df_from_sql_risk_reg = pd.read_sql('SELECT * FROM public."_6wsd21_edms_nec_risk_register";', conn)
        df_from_sql_risk_reg.rename(columns=mappings_risk_reg, inplace=True)

        # Get starting date
        start_date = df_all.sort_values(by="Doc_Date", ascending=True)["Doc_Date"].iloc[0]
        # Result Dataframe
        result_df = pd.DataFrame()

        for i in range(6):
            curr_date = update_months(today, -i)
            curr_month_string = curr_date.strftime('%Y%m')
            next_date = update_months(curr_date, 1)
            next_date_string = next_date.strftime('%Y-%m') + '-01'
            # Filter all the respective dates
            df_from_sql_nec = df_all[df_all['Doc_Date'] < next_date_string]

            # Initialize an empty DataFrame
            df = pd.DataFrame()

            # A1 - A3
            # Clean and convert the specific columns to float
            # A1
            if not df_from_sql_nec_section_of_works.empty:
                pwdd = df_from_sql_nec_section_of_works['Cumulative_PWDD'].str.replace(',', '').str.strip().astype(float)
                fcst_final_pwdd = df_from_sql_nec_section_of_works['Forecast_of_the_final_Prices_for_the_Work_Done_to_Date__PWDD'].str.replace(',', '').str.strip().astype(float)
                df['A1. PWDD'] = pwdd
                df['A1. Fcst_Final_PWDD'] = fcst_final_pwdd.round(2)
                df['A1. PWDD_to_Fcst_Final_PWDD'] = ((df['A1. PWDD']/df['A1. Fcst_Final_PWDD'])*100).round(2)
            else:
                pwdd = 0
                fcst_final_pwdd = 0
                df['A1. PWDD'] = 0
                df['A1. Fcst_Final_PWDD'] = 0
                df['A1. PWDD_to_Fcst_Final_PWDD'] = 100

            # A2
            if not df_from_sql_nec_section_of_works.empty:
                df['A2. Fcst_Final_PWDD'] = fcst_final_pwdd.round(2)
                fcst_final_total_prices = df_from_sql_nec_section_of_works['Latest_Forecast_Total_of_the_Prices'].str.replace(',', '').str.strip().astype(float)
                df['A2. Fcst_Final_Total_Prices'] =  fcst_final_total_prices.round(2)
                df['A2. Fcst_Final_PWDD_to_Fcst_Final_Total_Prices'] = ((df['A2. Fcst_Final_PWDD']/df['A2. Fcst_Final_Total_Prices'])*100).round(2)

                if ((fcst_final_pwdd/fcst_final_total_prices)*100 < 100).bool():
                    df['A2. Scenario'] = 'A'
                    df['A2. PainGain'] = (df['A2. Fcst_Final_Total_Prices'] - df['A2. Fcst_Final_PWDD'])*0.5
                elif ((fcst_final_pwdd/fcst_final_total_prices)*100 < 110).bool():
                    df['A2. Scenario'] = 'B'
                    df['A2. PainGain'] = (df['A2. Fcst_Final_Total_Prices'] - df['A2. Fcst_Final_PWDD'])*0.5
                else:
                    df['A2. Scenario'] = 'C'
                    df['A2. PainGain'] = ((df['A2. Fcst_Final_PWDD'] - (df['A2. Fcst_Final_PWDD']*1.1)) - df['A2. Fcst_Final_Total_Prices']*0.1*0.5).round(2)
            else:
                fcst_final_total_prices = 0
                df['A2. Fcst_Final_PWDD'] = 0
                df['A2. Fcst_Final_Total_Prices'] =  0
                df['A2. Fcst_Final_PWDD_to_Fcst_Final_Total_Prices'] = 100
                df['A2. Scenario'] = 'A'
                df['A2. PainGain'] = 0

            # A3
            if not df_from_sql_nec_section_of_works.empty:
                df['A3. Changed_Total_Price'] = abs(round((fcst_final_pwdd - fcst_final_total_prices), 2))
                df['A3. Fcst_Final_Total_Prices'] = round(fcst_final_total_prices, 2)
                df['A3. Changed_Total_Price_to_Fcst_Final_Total_Prices'] = abs(round((fcst_final_pwdd - fcst_final_total_prices)/fcst_final_total_prices * 100, 2))
            else:
                df['A3. Changed_Total_Price'] = 0
                df['A3. Fcst_Final_Total_Prices'] = 0
                df['A3. Changed_Total_Price_to_Fcst_Final_Total_Prices'] = 0
            
            # B1 - B5
            # Convert 'Revised_Completion_Date' to datetime, errors='coerce' will handle None and invalid dates
            if 'Ori_Completion_Date' in df_from_sql_nec and df_from_sql_nec['Ori_Completion_Date'].any():
                df_from_sql_nec['Ori_Completion_Date'] = pd.to_datetime(df_from_sql_nec['Ori_Completion_Date'], errors='coerce')
            else:
                df_from_sql_nec['Ori_Completion_Date'] = None
            
            if 'Revised_Completion_Date' in df_from_sql_nec and df_from_sql_nec['Revised_Completion_Date'].any():
                df_from_sql_nec['Revised_Completion_Date'] = pd.to_datetime(df_from_sql_nec['Revised_Completion_Date'], errors='coerce')
            else:
                df_from_sql_nec['Revised_Completion_Date'] = None
            
            # B1
            # Find the latest 'Revised_Completion_Date'
            if (not df_from_sql_nec['Revised_Completion_Date'].isnull().all()):
                latest_row = df_from_sql_nec.loc[df_from_sql_nec['Revised_Completion_Date'].idxmax()]
                # Extract the latest 'Revised_Completion_Date'
                latest_date = latest_row['Revised_Completion_Date']
            else:
                latest_date = None

            # Assuming df_from_sql_nec_section_of_works is already defined and contains 'starting_date'
            if not df_from_sql_nec_section_of_works.empty:
                df['contract_start_date'] = df_from_sql_nec_section_of_works['starting_date'].dt.tz_localize(None)
                # Calculate today's date
                today = pd.to_datetime(datetime.today())
                time_elapsed = (today - df['contract_start_date']).dt.days

                # Calculate the time elapsed from contract start date to today
                df['B1. Time_Elapsed'] = time_elapsed
                if (latest_date):
                    # Assign the latest 'Revised_Completion_Date' to 'Longest Section / Key day' in df
                    df['Longest Section / Key day'] = latest_date
                    # Calculate the total contractual duration
                    df['B1. Contractual_Duration'] = (df['Longest Section / Key day'].dt.tz_localize(None) - df['contract_start_date']).dt.days
                    # Calculate the ratio of time elapsed to contractual duration as a percentage
                    df['B1. Time_Elapsed_to_Contractual_Duration'] = ((df['B1. Time_Elapsed'] / df['B1. Contractual_Duration']) * 100).round(2)
                else:
                    df['Longest Section / Key day'] = None
                    df['B1. Contractual_Duration'] = 0
                    df['B1. Time_Elapsed_to_Contractual_Duration'] = 0
            else:
                df['contract_start_date'] = None
                df['B1. Time_Elapsed'] = 0
                df['Longest Section / Key day'] = None
                df['B1. Contractual_Duration'] = 0
                df['B1. Time_Elapsed_to_Contractual_Duration'] = 0

            # B2
            if 'Planned_Completion_Date_PCD' in df_from_sql_PCD.columns:
                df_from_sql_PCD['Planned_Completion_Date_PCD'] = pd.to_datetime(df_from_sql_PCD['Planned_Completion_Date_PCD'], errors="coerce")

                latest_row = df_from_sql_PCD.loc[df_from_sql_PCD['Planned_Completion_Date_PCD'].idxmax()]
                latest_planned_date = latest_row['Planned_Completion_Date_PCD']
                df['Latest Planned Date'] = latest_planned_date

                # Calculate the time elapsed from contract start date to today
                df['B2. Time_Elapsed'] = time_elapsed
                # Calculate the total planned duration for completion
                df['B2. Planned_Duration'] = (df['Latest Planned Date'].dt.tz_localize(None) - df['contract_start_date']).dt.days
                df['B2. Time_Elapsed_to_Planned_Duration'] = ((df['B2. Time_Elapsed'] / df['B2. Planned_Duration']) * 100).round(2)
            else:
                latest_row = None
                latest_planned_date = None
                df['B2. Time_Elapsed'] = 0
                df['B2. Planned_Duration'] = 0
                df['B2. Time_Elapsed_to_Planned_Duration'] = 0

            # B3
            df_json = pd.DataFrame()
            # Convert date columns to datetime objects
            if 'key_Date' in df_from_sql_nec.columns and df_from_sql_nec['key_Date'].any() and df_from_sql_nec['Revised_Completion_Date'] and df_from_sql_nec['Ori_Completion_Date']:
                df_json['Key_Date'] = df_from_sql_nec['key_Date']
                df_json['Revised_Completion_Date'] = pd.to_datetime(df_from_sql_nec['Revised_Completion_Date'])
                df_json['Ori_Completion_Date'] = pd.to_datetime(df_from_sql_nec['Ori_Completion_Date'])
                
                # Calculate the latest "Revised_Completion_Date" for each section or key date
                latest_revised_completion = df_json.groupby(['Key_Date'])['Revised_Completion_Date'].max().reset_index()
                latest_revised_completion.rename(columns={'Revised_Completion_Date': 'Section or Key day Revised_Completion_Date'}, inplace=True)
                
                # Calculate the earliest "Ori_Completion_Date" for each section or key date
                earliest_ori_completion = df_json.groupby(['Key_Date'])['Ori_Completion_Date'].min().reset_index()
                earliest_ori_completion.rename(columns={'Ori_Completion_Date': 'Section or Key day Ori_Completion_Date'}, inplace=True)
                
                # Merge the latest and earliest dates into a single DataFrame
                merged_dates = pd.merge(latest_revised_completion, earliest_ori_completion, on=['Key_Date'])
                
                # Calculate the EOT (Extension Of Time) in days
                merged_dates['EOT'] = (merged_dates['Section or Key day Revised_Completion_Date'] - merged_dates['Section or Key day Ori_Completion_Date']).dt.days
                merged_dates = merged_dates[merged_dates['Key_Date'].notnull() & (merged_dates['Key_Date'] != '')]
                # merged_dates = merged_dates.set_index('Key_Date')

                if not merged_dates['Section or Key day Revised_Completion_Date'].isnull().all() and not merged_dates['Section or Key day Revised_Completion_Date'].isnull().all():
                    latest_row = merged_dates.loc[merged_dates['Section or Key day Revised_Completion_Date'].idxmax()]
                    # For the longest section, excluding establishment works
                    df['B3. Extended_Completion_Date'] = latest_row['Section or Key day Revised_Completion_Date'].strftime('%Y%m%d')
                    df['B3. Original_Completion_Date'] = latest_row['Section or Key day Ori_Completion_Date'].strftime('%Y%m%d')
                    df['B3. Extension_of_Time_of_Contract'] = latest_row['EOT']
                else:
                    df['B3. Extended_Completion_Date'] = None
                    df['B3. Original_Completion_Date'] = None
                    df['B3. Extension_of_Time_of_Contract'] = None
            else:
                merged_dates = pd.DataFrame()
                df['B3. Extended_Completion_Date'] = None
                df['B3. Original_Completion_Date'] = None
                df['B3. Extension_of_Time_of_Contract'] = None
            
            # B4
            def generate_key_date(row):
                return {
                    'ID': row['key_Date'],
                    'Contract_Date': row['Section or Key day Ori_Completion_Date'],
                    'Planned_Date': row['Planned_Completion_Date_PCD'],
                    'Updated_Date': row['Section or Key day Revised_Completion_Date']
                }
            
            if 'Key_Date' in merged_dates.columns and 'key_Date' in df_from_sql_PCD.columns:
                filtered_merged_dates_key = merged_dates[merged_dates['Key_Date'].str.lower().str.strip().str.startswith('key date')]
                merged_PCD = pd.merge(df_from_sql_PCD, filtered_merged_dates_key, left_on='key_Date', right_on='Key_Date', how='inner') 
                merged_PCD = merged_PCD[['key_Date', 'Section or Key day Revised_Completion_Date', 'Section or Key day Ori_Completion_Date', 'Planned_Completion_Date_PCD']]
                merged_PCD['Section or Key day Ori_Completion_Date'] = pd.to_datetime(merged_PCD['Section or Key day Ori_Completion_Date']).dt.strftime('%Y%m%d')
                merged_PCD['Section or Key day Revised_Completion_Date'] = pd.to_datetime(merged_PCD['Section or Key day Revised_Completion_Date']).dt.strftime('%Y%m%d')
                merged_PCD['Planned_Completion_Date_PCD'] = pd.to_datetime(merged_PCD['Planned_Completion_Date_PCD']).dt.strftime('%Y%m%d')
                key_dates_json = merged_PCD.apply(lambda row: generate_key_date(row), axis=1).to_json(orient='records')
                df['B4. KEY_DATES'] = '{ "KEY_DATE": ' + key_dates_json + '}'
            else:
                df['B4. KEY_DATES'] = None

            # B5
            if 'Key_Date' in merged_dates.columns and 'key_Date' in df_from_sql_PCD.columns:
                filtered_merged_dates_section = merged_dates[merged_dates['Key_Date'].str.strip().str.lower().str.startswith('section')]
                merged_PCD = pd.merge(df_from_sql_PCD, filtered_merged_dates_section, left_on='key_Date', right_on='Key_Date', how='inner') 
                merged_PCD = merged_PCD[['key_Date', 'Section or Key day Revised_Completion_Date', 'Section or Key day Ori_Completion_Date', 'Planned_Completion_Date_PCD']]
                merged_PCD['Section or Key day Ori_Completion_Date'] = pd.to_datetime(merged_PCD['Section or Key day Ori_Completion_Date']).dt.strftime('%Y%m%d')
                merged_PCD['Section or Key day Revised_Completion_Date'] = pd.to_datetime(merged_PCD['Section or Key day Revised_Completion_Date']).dt.strftime('%Y%m%d')
                merged_PCD['Planned_Completion_Date_PCD'] = pd.to_datetime(merged_PCD['Planned_Completion_Date_PCD']).dt.strftime('%Y%m%d')
                sec_dates_json = merged_PCD.apply(lambda row: generate_key_date(row), axis=1).to_json(orient='records')
                df['B5. SEC_DATES'] = '{ "SEC_DATE": ' + sec_dates_json + '}'
            else:
                df['B5. SEC_DATES'] = None
            
            # Section C1 - C5
            # C1
            # filtered_df_CEW = df_from_sql_nec[
            #     (df_from_sql_nec['NEC_Doc_Type'] == 'EW-') &
            #     (df_from_sql_nec['From'].str.startswith('DCK JV')) &
            #     (df_from_sql_nec['Doc_Ver'] == '0') | (df_from_sql_nec['Doc_Ver'] == 0)
            # ]
            
            filtered_df_EW = df_from_sql_risk_reg[~df_from_sql_risk_reg['Notification_Reference'].isnull() & ~df_from_sql_risk_reg['Date_of_Early_Warning'].isnull()]
            # Get the total number of records that meet the conditions
            filtered_df_EW_total_records = len(filtered_df_EW)
            df['C3. Total_Num_EW'] = filtered_df_EW_total_records

            filtered_df_CEWN = filtered_df_EW[filtered_df_EW['Notified_by'].str.upper() == 'C']
            filtered_df_PM = filtered_df_EW[filtered_df_EW['Notification_Reference'].str.upper() == 'PM']

            df['C1. Total_Num_EW_by_Contractor'] = len(filtered_df_CEWN)
            df['C2. Total_Num_EW_by_Project_Manager'] = len(filtered_df_PM)
            
            # C4
            # Convert the relevant columns to datetime objects, handling errors
            if not df_from_sql_risk_reg['Date_of_Close_of_EW'].isnull().all():
                df_from_sql_risk_reg['Date_of_Close_of_EW'] = pd.to_datetime(df_from_sql_risk_reg['Date_of_Close_of_EW'], errors='coerce')  
                filtered_df_Closed_EW = df_from_sql_risk_reg[df_from_sql_risk_reg['Status'] == 'Closed']
            else:
                filtered_df_Closed_EW = pd.DataFrame()

            df_from_sql_risk_reg['Date_of_Early_Warning'] = pd.to_datetime(df_from_sql_risk_reg['Date_of_Early_Warning'], errors='coerce')
            df['C4. Closed_EW'] = len(filtered_df_Closed_EW)
            df['C4. Total_EW'] = len(df_from_sql_risk_reg)
            if len(df_from_sql_risk_reg) == len(filtered_df_Closed_EW) :
                df['C4. Resolved_EW_To_Total_EW'] = 100
            else:
                df['C4. Resolved_EW_To_Total_EW'] = round(((len(filtered_df_Closed_EW) / len(df_from_sql_risk_reg))*100),2)

            # C5
            # Filter out rows where either date is null
            filtered_rr_df = df_from_sql_risk_reg.dropna(subset=['Date_of_Close_of_EW'])
            filtered_rr_df.dropna(subset=['Date_of_Early_Warning'], inplace=True)

            if not filtered_rr_df.empty:
                # Calculate the difference in days between Date_of_Close_of_EW and Date_of_Early_Warning
                filtered_rr_df['Duration_Days'] = (filtered_rr_df['Date_of_Close_of_EW'] - filtered_rr_df['Date_of_Early_Warning']).dt.days

                duration_json = filtered_rr_df['Duration_Days'].to_json(orient='records')
                df['C5. Durations_All_Closed_EW'] = '{ "Duration": ' + duration_json + '}'
                df['C5. Num_closed_EW'] = len(filtered_rr_df)
                # Calculate the average duration
                average_duration = round(filtered_rr_df['Duration_Days'].mean(),2)
                df['C5. Avg_Duration_to_Resolve_EW'] = average_duration
            else:
                df['C5. Durations_All_Closed_EW'] = None
                df['C5. Num_closed_EW'] = 0
                df['C5. Avg_Duration_to_Resolve_EW'] = 0
            
            # Filtered dataframe for difference event types
            # filtered_df_PMI= df_from_sql_nec[(df_from_sql_nec['NEC_Doc_Type'] == 'PMI-') &
            #     (df_from_sql_nec['Doc_Ver'] == '0') | (df_from_sql_nec['Doc_Ver'] == 0)
            # ]
            # filtered_df_NCE= df_from_sql_nec[(df_from_sql_nec['NEC_Doc_Type'] == 'NCE-') &
            #     (df_from_sql_nec['Doc_Ver'] == '0') | (df_from_sql_nec['Doc_Ver'] == 0)
            # ]
            # filtered_df_PMN= df_from_sql_nec[(df_from_sql_nec['NEC_Doc_Type'] == 'PMN-') &
            #     (df_from_sql_nec['Doc_Ver'] == '0') | (df_from_sql_nec['Doc_Ver'] == 0)
            # ]
            # filtered_df_QA= df_from_sql_nec[((df_from_sql_nec['NEC_Doc_Type'] == 'QA-') & 
            #     (df_from_sql_nec['Doc_Ver'] == '0') | (df_from_sql_nec['Doc_Ver'] == 0)) & 
            #     ~df_from_sql_nec['From_Status'].isnull()
            # ]

            filtered_df_PMI = df_p2[df_p2['Doc_Date'] < next_date_string]
            filtered_df_NCE = df_c2[df_c2['Doc_Date'] < next_date_string]
            # filtered_df_PMN = df_p3[df_p3['Doc_Date'] < next_date_string]
            filtered_df_QA = df_p4[df_p4['Doc_Date'] < next_date_string]
            filtered_df_PMN = pd.DataFrame()

            # D1, D2, D3
            df['D1. PM_Instruction']=len(filtered_df_PMI)
            df['D2. Contractor_NCE']=len(filtered_df_NCE)
            df['D3. PM_NCE']=len(filtered_df_PMN)
            
            # D3a
            # Initialize a counter for accepted records and non-accepted records
            accepted_count = 0
            # Iterate through each filtered NCE record
            for _, row in filtered_df_NCE.iterrows():
                nec_event_no = row['NEC_Event_No']
                # Check if there's a 'PMN-' record with the same NEC_Event_No
                if not df_from_sql_nec[df_from_sql_nec['NEC_Event_No'] == nec_event_no].empty:
                    accepted_count += 1
            df['D3a. Contractor_NCE_PM_decision'] = len(filtered_df_NCE)
            df['D3a. Contractor_NCE_PM_accepted_instructed'] = accepted_count
            df['D3a. Ratio'] = round(accepted_count / len(filtered_df_NCE), 3)

            # D4
            df['D4. Total_NCE']= len(filtered_df_NCE) + len(filtered_df_PMN)
            
            # D5
            # Filter records where NEC_Clause starts with '60.1'
            filtered_ground_clause_df = df_from_sql_nec[df_from_sql_nec['NEC_Clause'].str.contains('60.1')]
            # Function to extract the classification of ground from NEC_Clause
            def extract_classification(nec_clause):
                match = re.search(r'60\.1\((\d+)\)', nec_clause)
                if match:
                    return int(match.group(1))
                return None
            def generate_compensation_events(row):
                return pd.Series({
                    'CE_ID': row['CE_No'],
                    'Ground_ID': row['NEC_Clause'].replace('60.1', '').replace('(', '').replace(')', '')
                })

            # Apply the function to the filtered DataFrame
            filtered_ground_clause_df['Classification_of_Ground'] = filtered_ground_clause_df['NEC_Clause'].apply(extract_classification)
            # Drop rows where classification extraction failed (if any)
            filtered_ground_clause_df = filtered_ground_clause_df.dropna(subset=['Classification_of_Ground'])
            filtered_ground_clause_df_json = filtered_ground_clause_df.apply(lambda row: generate_compensation_events(row), axis=1).to_json(orient='records')
            
            df['D5. Total_CE'] = len(filtered_ground_clause_df)
            df['D5. CEs'] = '{ "CE": ' + filtered_ground_clause_df_json + '}'
            
            # D6
            df['D6. Num_Implemented_Events'] = len(filtered_df_QA)
            df['D6. Num_Notified_Events'] = len(filtered_df_PMN)
            if len(filtered_df_QA) == len(filtered_df_PMN):
                df['D6. Ratio']= 100
            elif len(filtered_df_PMN):
                df['D6. Ratio']= round((len(filtered_df_QA) / (len(filtered_df_PMN)))*100,2)
            else:
                df['D6. Ratio']= 0

            def calculate_date_difference(group, table_name_1, table_name_2, table_key_1, table_key_2):
                date_1 = group.loc[group['Table'] == table_name_1, 'Doc_Date']
                date_2 = group.loc[group['Table'] == table_name_2, 'Doc_Date']
                if not date_1.empty and not date_2.empty:
                    index_1 = len(date_1)-1
                    index_2 = len(date_2)-1
                    # Calculate the difference in days
                    date_diff = (date_2.iloc[index_2] - date_1.iloc[index_1]).days
                    return pd.Series({
                        'NEC_Event_No': group['NEC_Event_No'].iloc[0],
                        'CE_No': group['CE_No'].iloc[0],
                        table_key_1: date_1.iloc[index_1].strftime('%Y%m%d'),
                        table_key_2: date_2.iloc[index_2].strftime('%Y%m%d'),
                        'Duration': date_diff,
                    })
                return None

            # D7
            # Group by NEC_Event_No and calculate the date difference
            # def calculate_date_difference(group, pmn_table_name='p3', ie_table_name='p4'):
            #     pmn_date = group.loc[group['Table'] == pmn_table_name, 'Doc_Date']
            #     ie_date = group.loc[group['Table'] == ie_table_name, 'Doc_Date']
            #     if not pmn_date.empty and not qa_date.empty:
            #         pmn_index = len(pmn_date)-1
            #         ie_index = len(ie_date)-1
            #         # Calculate the difference in days
            #         date_diff = (ie_date.iloc[ie_index] - pmn_date.iloc[pmn_index]).days
            #         return pd.Series({
            #             'NEC_Event_No': group['NEC_Event_No'].iloc[0],
            #             'CE_No': group['CE_No'].iloc[0],
            #             'Date_Notification': pmn_date.iloc[pmn_index].strftime('%Y%m%d'),
            #             'Date_Implementation': ie_date.iloc[ie_index].strftime('%Y%m%d'),
            #             'Duration': date_diff
            #         })
            #     return None

            # Apply the calculation to each group and filter out None results
            # Apply the calculation to each group using a lambda function to pass parameters
            NCE_QA_date_diff_df = df_from_sql_nec.groupby('NEC_Event_No').apply(lambda group: calculate_date_difference(group, 'p3', 'p4', 'Date_Notification', 'Date_Implementation')).dropna().reset_index(drop=True)
            
            if not NCE_QA_date_diff_df.empty:
                implemented_ces = NCE_QA_date_diff_df[['Date_Notification', 'Date_Implementation', 'Duration']].to_json(orient='records')
                df['D7. Implemented_CEs'] = '{ "Implemented_CE": ' + implemented_ces + '}'
                df['D7. Num_implemented'] = len(NCE_QA_date_diff_df)
                df['D7. Average_duration']= round(NCE_QA_date_diff_df['Duration'].mean(), 2)
            else:
                df['D7. Implemented_CEs'] = None
                df['D7. Num_implemented'] = 0
                df['D7. Average_duration']= 0
            
            # D8
            # def calculate_date_difference(group, pmn_table_name='p3', csq_doc_type='p4'):
            #     pmn_date = group.loc[group['NEC_Doc_Type'] == pmn_doc_type, 'Doc_Date']
            #     csq_date = group.loc[group['NEC_Doc_Type'] == csq_doc_type, 'Doc_Date']
            #     if not pmn_date.empty and not csq_date.empty:
            #         pmn_index = len(pmn_date)-1
            #         csq_index = len(csq_date)-1               
            #         # Calculate the difference in days
            #         date_diff = (csq_date.iloc[csq_index] - pmn_date.iloc[pmn_index]).days
            #         return pd.Series({
            #             'NEC_Event_No': group['NEC_Event_No'].iloc[0],
            #             'CE_No': group['CE_No'].iloc[0],
            #             'Date_Quotation_Req': pmn_date.iloc[pmn_index].strftime('%Y%m%d'),
            #             'Date_Quotation_Sub': csq_date.iloc[csq_index].strftime('%Y%m%d'),
            #             'Duration': date_diff
            #         })
            #     return None

            PMN_CQS_date_diff_df = df_from_sql_nec.groupby('NEC_Event_No').apply(lambda group: calculate_date_difference(group, 'p3', 'p4', 'Date_Quotation_Req', 'Date_Quotation_Sub')).dropna().reset_index(drop=True)
            
            if not PMN_CQS_date_diff_df.empty:
                quotation_subs = PMN_CQS_date_diff_df[['Date_Quotation_Req', 'Date_Quotation_Sub', 'Duration']].to_json(orient="records")
                df['D8. Quotation_Subs'] = '{ "Quotation_Sub": ' + quotation_subs + '}'
                df['D8. Num_quotation'] = len(PMN_CQS_date_diff_df)
                df['D8. Average_duration'] = round(PMN_CQS_date_diff_df['Duration'].mean(), 2)
            else:
                df['D8. Quotation_Subs'] = None
                df['D8. Num_quotation'] = 0
                df['D8. Average_duration'] = 0
            
            # D9
            # def calculate_date_difference(group, csq_doc_type='CSQ-', qa_doc_type='QA-'):
            #     csq_date = group.loc[group['NEC_Doc_Type'] == csq_doc_type, 'Doc_Date']
            #     qa_date = group.loc[group['NEC_Doc_Type'] == qa_doc_type, 'Doc_Date']

            #     if not csq_date.empty and not qa_date.empty:
            #         qa_index = len(qa_date)-1
            #         csq_index = len(csq_date)-1
            #         # Calculate the difference in days
            #         date_diff = (qa_date.iloc[qa_index] - csq_date.iloc[csq_index]).days
            #         return pd.Series({
            #             'NEC_Event_No': group['NEC_Event_No'].iloc[0],
            #             'CE_No': group['CE_No'].iloc[0],
            #             'Date_Quotation_Submission': csq_date.iloc[csq_index].strftime('%Y%m%d'),
            #             'Date_PM_Reponse': qa_date.iloc[qa_index].strftime('%Y%m%d'),
            #             'Duration': date_diff
            #         })
            #     return None
            
            CSQ_QA_date_diff_df = df_from_sql_nec.groupby('NEC_Event_No').apply(lambda group: calculate_date_difference(group, 'c3', 'p4', 'Date_Quotation_Submission', 'Date_PM_Reponse')).dropna().reset_index(drop=True)

            if not CSQ_QA_date_diff_df.empty:
                quotation_assessments = CSQ_QA_date_diff_df[['Date_Quotation_Submission', 'Date_PM_Reponse', 'Duration']].to_json(orient='records')
                df['D9. Quotation_Assessments'] = '{ "Quotation_Assessment": ' + quotation_assessments + '}'
                df['D9. Num_PM_Response'] = len(CSQ_QA_date_diff_df)
                df['D9. Average_duration'] = round(CSQ_QA_date_diff_df['Duration'].mean(), 2)
            else:
                df['D9. Quotation_Assessments'] = None
                df['D9. Num_PM_Response'] = 0
                df['D9. Average_duration'] = 0  
            
            # D10 - D11
            # Calculate the Cost Implication
            def calculate_cost(row):
                if row['CE_Increase_Decrease'].lower() == 'increase':
                    return row['CE_Amount']
                elif row['CE_Increase_Decrease'].lower() == 'decrease':
                    return -row['CE_Amount']
                else: 
                    return 0
            def generate_implemented_compensations(row):
                if 'Extension_in_days' not in row:
                    return pd.Series({
                        'Cost_Implication': round(row['Cost_Implication'], 2),
                        'Time_Implication': 0
                    })
                return pd.Series({
                    'Cost_Implication': round(row['Cost_Implication'], 2),
                    'Time_Implication': row['Extension_in_days']
                })
            
            if 'CE_Amount' in filtered_df_QA.columns or 'Change_to_Time' in filtered_df_QA.columns:
                # Drop rows where CE_PMI_Amount or Extension_in_days is empty
                filtered_pmi_amount = filtered_df_QA.dropna(subset=['CE_Amount'])

                if not filtered_pmi_amount.empty:
                    filtered_pmi_amount['Cost_Implication'] = filtered_pmi_amount.apply(calculate_cost, axis=1)
                    df_groupby_ori_doc_no = filtered_pmi_amount.groupby(['CE_No']).sum()

                    implemented_compensations = df_groupby_ori_doc_no.apply(lambda row: generate_implemented_compensations(row), axis=1).dropna().reset_index(drop=True)
                    implemented_compensations_json = implemented_compensations.to_json(orient='records')
                    df['D10. Implemented_Compensations'] = '{ "Implemented_Compensation": ' + implemented_compensations_json + '}'

                    total_cost_implication = filtered_pmi_amount['Cost_Implication'].sum()
                    df['D10. Sum_Cost_Implication'] = round(total_cost_implication, 2)
                    df['D10a. Avg_Cost_Implication'] = round((total_cost_implication / len(filtered_pmi_amount)), 2)

                    df['D11. Time_Cost_Implication'] = 0

                else:
                    df['D10. Implemented_Compensations'] = None
                    df['D10. Sum_Cost_Implication'] = 0
                    df['D10a. Avg_Cost_Implication'] = 0
                    df['D11. Time_Cost_Implication'] = 0
            else:
                df['D10. Implemented_Compensations'] = None
                df['D10. Sum_Cost_Implication'] = 0
                df['D10a. Avg_Cost_Implication'] = 0
                df['D11. Time_Cost_Implication'] = 0

            
            # Include starting date
            df['start_date'] = start_date
            # Include the year month
            df['year_month'] = curr_month_string

            print('DataFrame:', df)
            df.columns = df.columns.str.replace(' ', '_').str.replace('.', '_').str.replace('(', '_').str.replace(')', '').str.replace('%', 'percent')
            result_df = result_df.append(df)

        # Write the DataFrame back to a SQL table
        result_df.to_sql('nec_6wsd21_icwps', con=conn, if_exists='replace', index=False)

# */2 * * * * Execute every two minute 
with DAG(
        dag_id="1nec_6wsd21_icwps",
        schedule_interval="0 1,5,9,12,17 * * *",
        default_args={
            "owner": "airflow",
            "retries": 1,
            "retry_delay": timedelta(minutes=5),
            "start_date": datetime(2022, 10, 24)
        },
        catchup=False) as f:
    
    getDataAndSendToPSQL = PythonOperator(
        task_id="getDataAndSendToPSQL",
        python_callable=getMongoDB,
        op_kwargs={"name": "Dylan"},
        provide_context=True,
    )

    # reformData = PythonOperator(
    #     task_id="reformData",
    #     python_callable=reformData,
    #     provide_context=True,
    #     # op_kwargs={"name": "Dylan"}
    # )

    getDrowToken = PythonOperator(
        task_id="getDrowToken",
        python_callable=getDrowToken,
        provide_context=True,
        # op_kwargs={"name": "Dylan"}
    )

getDrowToken >> getDataAndSendToPSQL