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

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()

    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
        df_all = pd.read_sql("""SELECT * FROM public.nec_c04 
                                         WHERE "Doc_Date" IS NOT NULL 
                                         ;""", conn)
        df_from_sql_nec_section_of_works = pd.read_sql("SELECT * FROM public.nec_section_of_work", conn)
        df_from_sql_PCD = pd.read_sql("SELECT * FROM public.c04_key_date_data;", conn)
        df_from_sql_risk_reg = pd.read_sql("SELECT * FROM public.nec_risk_register;", conn)

        # 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()
            df["A0. contract no"] = ["ND/2019/04"]
            df["A0. contract type"] = "NEC3"
            df["A0. contract form"] = "ECC"
            df["A0. main option"] = "C"
            df["A0. secondary clauses"] = "X1, X5, X7, X14, X15, X16, X20, Z"
            df["A0. description of works"] = """
        The works are the construction of viaduct, underpass and depressed
        road structure, footbridge and associated lift and stairs, stormwater
        pumping station, sewage pumping station, public toilet, refuse
        collection point facility, at-grade road and associated slope / retaining
        wall, noise barrier and semi-enclosure, junction improvement works
        and landscaping works as more particularly described in the General
        Particulars Clause 2.1 under the Works Information.
        """
            df["A0. starting date"]= pd.to_datetime("2020-08-14").strftime('%Y%m%d')
            df["A0. completion date"] = pd.to_datetime("2025-07-09").strftime('%Y%m%d')
            df["A0. tendered prices"] =  2148000000
            df["A0. name of contractor"] = "DCK JV (Daewoo - Chun Wo - Kwan Lee Joint Venture)"
            df["A0. name of consultant"] = "AECOM Asia Company Limited"        
            # A1 - A3
            # Clean and convert the specific columns to float
            # A1
            pwdd = df_from_sql_nec_section_of_works['Cumulative_PWDD'].str.replace(',', '').str.strip().astype(float)
            df['A1. PWDD'] = pwdd
            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. 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)

            # A2
            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)

            # A3
            df['A3. Changed_Total_Price'] = 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'] = round((fcst_final_pwdd - fcst_final_total_prices)/fcst_final_total_prices * 100, 2)
            
            # B1 - B5
            # Convert 'Revised_Completion_Date' to datetime, errors='coerce' will handle None and invalid dates
            df_from_sql_nec['Ori_Completion_Date'] = pd.to_datetime(df_from_sql_nec['Ori_Completion_Date'], errors='coerce')
            df_from_sql_nec['Revised_Completion_Date'] = pd.to_datetime(df_from_sql_nec['Revised_Completion_Date'], errors='coerce')

            # B1
            # Find the latest 'Revised_Completion_Date'
            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']
            # Assuming df_from_sql_nec_section_of_works is already defined and contains 'starting_date'
            df['contract_start_date'] = df_from_sql_nec_section_of_works['starting_date'].dt.tz_localize(None)
            # Assign the latest 'Revised_Completion_Date' to 'Longest Section / Key day' in df
            df['Longest Section / Key day'] = latest_date
            # 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
            # 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)
            
            # B2
            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)

            # B3
            df_json = pd.DataFrame()
            # Convert date columns to datetime objects
            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
            
            # 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']
                }
            
            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 + '}'

            # B5
            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 + '}'
            
            # 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['EW_ref_'].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['EW_ref_'].str.upper().str.contains('CEWN')]
            filtered_df_PM = filtered_df_EW[filtered_df_EW['EW_ref_'].str.upper().str.contains('PM')]

            df['C1. Total_Num_EW_by_Contractor'] = len(filtered_df_CEWN)
            df['C2. Total_Num_EW_by_Project_Manager'] = len(filtered_df_PM)
            
            # filtered_df_PMEW = 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)
            # ]
            # Get the total number of records that meet the conditions
            # filtered_df_PMEW_total_records = len(filtered_df_PMEW)
            # df['C2. Total_Num_EW_by_Project_Manager'] = filtered_df_PMEW_total_records
            # df['C3. Total_Num_EW'] = filtered_df_CEW_total_records + filtered_df_PMEW_total_records
            
            # C4
            # Convert the relevant columns to datetime objects, handling errors
            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')
            df_from_sql_risk_reg['Date_of_Early_Warning'] = pd.to_datetime(df_from_sql_risk_reg['Date_of_Early_Warning'], errors='coerce')
            filtered_df_Closed_EW = df_from_sql_risk_reg[df_from_sql_risk_reg['Status'] == 'Close']
            
            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', 'Date_of_Early_Warning'])
            # 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
            
            # 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()
            ]

            # 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) & (df_from_sql_nec['NEC_Doc_Type'] == 'PMN-')].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.startswith('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['Original_Doc_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
            else:
                df['D6. Ratio']= round((len(filtered_df_QA) / (len(filtered_df_PMN)))*100,2)

            # D7
            # Group by NEC_Event_No and calculate the date difference
            def calculate_date_difference(group, pmn_doc_type='PMN-', qa_doc_type='QA-'):
                pmn_date = group.loc[group['NEC_Doc_Type'] == pmn_doc_type, 'Doc_Date']
                qa_date = group.loc[group['NEC_Doc_Type'] == qa_doc_type, 'Doc_Date']
                if not pmn_date.empty and not qa_date.empty:
                    pmn_index = len(pmn_date)-1
                    qa_index = len(qa_date)-1
                    # Calculate the difference in days
                    date_diff = (qa_date.iloc[qa_index] - pmn_date.iloc[pmn_index]).days
                    return pd.Series({
                        'NEC_Event_No': group['NEC_Event_No'].iloc[0],
                        'CE_No': group['Original_Doc_No'].iloc[0],
                        'Date_Notification': pmn_date.iloc[pmn_index].strftime('%Y%m%d'),
                        'Date_Implementation': qa_date.iloc[qa_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, 'PMN-', 'QA-')).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
            filtered_df_CSQ= df_from_sql_nec[(df_from_sql_nec['NEC_Doc_Type'] == 'CSQ-') &
                (df_from_sql_nec['Doc_Ver'] == '0') |df_from_sql_nec['Doc_Ver'] == 0
            ]
            def calculate_date_difference(group, pmn_doc_type='PMN-', csq_doc_type='CSQ-'):
                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['Original_Doc_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, 'PMN-', 'CSQ-')).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['Original_Doc_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, 'CSQ-', 'QA-')).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_PMI_Amount']
                elif row['CE_Increase_Decrease'].lower() == 'decrease':
                    return -row['CE_PMI_Amount']
                else: 
                    return 0
            def generate_implemented_compensations(row):
                return pd.Series({
                    'Cost_Implication': round(row['Cost_Implication'], 2),
                    'Time_Implication': row['Extension_in_days']
                })
            
            if 'CE_PMI_Amount' in filtered_df_QA.columns or 'Extension_in_days' 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_PMI_Amount', 'Extension_in_days'])

                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(['Original_Doc_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'] = filtered_pmi_amount['Extension_in_days'].sum()

                    # # Drop rows where 'Revised_Completion_Date' or 'Ori_Completion_Date' is empty
                    # filtered_df_QA = filtered_df_QA.dropna(subset=['Revised_Completion_Date', 'Ori_Completion_Date'])
                    # filtered_df_QA['Cost_Implication'] = filtered_df_QA.apply(calculate_cost, axis=1)
                    # # Convert the date columns to datetime objects
                    # filtered_df_QA['Revised_Completion_Date'] = pd.to_datetime(filtered_df_QA['Revised_Completion_Date'], errors="coerce")
                    # filtered_df_QA['Ori_Completion_Date'] = pd.to_datetime(filtered_df_QA['Ori_Completion_Date'], errors="coerce")
                    # filtered_df_QA['Time_Implication'] = (filtered_df_QA['Revised_Completion_Date'] - filtered_df_QA['Ori_Completion_Date']).dt.days

                    # implemented_compensations = filtered_df_QA.apply(lambda row: generate_implemented_compensations(row), axis=1).dropna().reset_index(drop=True)
                    # df['D10. Implemented_Compensations'] = implemented_compensations.to_json(orient="records")

                    # total_cost_implication = filtered_df_QA['Cost_Implication'].sum()
                    # filtered_df_QA_non_zero = filtered_df_QA[(filtered_df_QA['CE_PMI_Amount'] != 0)]
                    # df['D10. Sum_Cost_Implication'] = round(total_cost_implication, 2)
                    # df['D10a. Avg_Cost_Implication'] =  round((total_cost_implication / len(filtered_df_QA_non_zero)), 2)
                
                    # total_time_implication = filtered_df_QA['Time_Implication'].sum()
                    # df['D11. Time Implication of Implemented Compensation Events'] = total_time_implication 
            
            # Include starting date
            df['start_date'] = start_date
            # Include the year month
            df['year_month'] = curr_month_string

            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_c04_icwp_v1', con=conn, if_exists='replace', index=False)

# */2 * * * * Execute every two minute 
with DAG(
        dag_id="1nec_c4_icwps_v1",
        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