doczyai-pipelines/fieldExtraction/tests/qa_qc.py

import pandas as pd
import os
import sys
from datetime import datetime
import re
from openpyxl import Workbook
from openpyxl.utils.dataframe import dataframe_to_rows
import argparse
import boto3
from botocore.exceptions import ClientError
sys.path.append("src")
import utils

file_name_column = "Contract Name"

"""python scripts/qa_qc.py --input_dir ______ --output_dir _______ --ac_df ________ --b_df ______
You should have both input and output (ouput is just an empty folder) and then one of ac_df or b_df or both
The final report will be saved in the parent directory of when you are running under reports/
Run this from the base folder of the repo."""


def parse_arguments():
    parser = argparse.ArgumentParser(description="QC/QA Integration Script")
    parser.add_argument("--input_dir", help="Input directory (local path or S3 URI)")
    parser.add_argument("--output_dir", help="Output directory")
    parser.add_argument("--ac_df", help="Path to AC DataFrame CSV")
    parser.add_argument("--b_df", help="Path to B DataFrame CSV")
    return parser.parse_args()


def is_s3_path(path):
    return path.startswith("s3://")


def read_s3_file(bucket, key):
    s3 = boto3.client("s3")
    try:
        response = s3.get_object(Bucket=bucket, Key=key)
        return response["Body"].read().decode("utf-8")
    except ClientError as e:
        print(f"Error reading S3 file {key} from bucket {bucket}: {e}")
        return None


def read_input_files(input_dir):
    if not input_dir:
        return {}

    if is_s3_path(input_dir):
        s3_parts = input_dir[5:].split("/", 1)
        bucket = s3_parts[0]
        prefix = s3_parts[1] if len(s3_parts) > 1 else ""
        s3 = boto3.client("s3")
        input_dict = {}
        try:
            for obj in s3.list_objects_v2(Bucket=bucket, Prefix=prefix)["Contents"]:
                key = obj["Key"]
                if key.endswith(".txt"):
                    content = read_s3_file(bucket, key)
                    if content:
                        input_dict[os.path.basename(key)] = content
        except ClientError as e:
            print(f"Error listing S3 objects: {e}")
    else:
        input_dict = {}
        for filename in os.listdir(input_dir):
            if filename.endswith(".txt"):
                with open(
                    os.path.join(input_dir, filename), "r", encoding="utf-8"
                ) as file:
                    input_dict[filename] = file.read()
    return input_dict


def remove_txt_extension(filename):
    if isinstance(filename, str) and filename.lower().endswith(".txt"):
        return filename[:-4]
    return filename


def detect_anomalies(series, field_name, is_merged=False):
    total_count = len(series)
    unique_values = series.fillna("").nunique()

    date_fields = [
        "CONTRACT_BASE_EFFECTIVE_DT",
        "CONTRACT_EFFECTIVE_DT",
        "CONTRACT_EFFECTIVE_DT_SIGNATORY",
        "CONTRACT_TERMINATION_DT",
        "CONTRACT_SIGNATORY_DT",
    ]
    date_pattern = r"(\d{1,2}[/-]\d{1,2}[/-]\d{2,4}|\d{2,4}[/-]\d{1,2}[/-]\d{1,2})"

    date_like_count = 0

    if pd.api.types.is_string_dtype(series) or pd.api.types.is_object_dtype(series):
        series = series.astype(str)
        if field_name in date_fields:
            date_like_count = series.str.count(date_pattern).sum()
        else:
            date_like_count = series.str.count(r"\d{1,2}/\d{1,2}/\d{2,4}").sum()

    data_type = "float64" if pd.api.types.is_float_dtype(series) else "string"

    dataset = "Merged DataFrame" if is_merged else "Original DataFrame"
    result = {
        f"Total entries ({dataset})": total_count,
        f"Unique values ({dataset})": unique_values,
        f"Date-like entries ({dataset})": date_like_count,
    }

    if is_merged:
        result[f"Data type ({dataset})"] = data_type

    return result


def perform_qc_qa_tests_both(ac_df, b_df, merged_df, input_dict, processed_files):
    total_input_files = len(input_dict) if input_dict else 0
    processed_ac_files = len(
        [f for f in processed_files if f in ac_df[file_name_column].values]
    )
    processed_b_files = len(
        [f for f in processed_files if f in b_df[file_name_column].values]
    )
    failed_ac_files = (
        total_input_files - processed_ac_files if total_input_files > 0 else 0
    )
    failed_b_files = (
        total_input_files - processed_b_files if total_input_files > 0 else 0
    )

    general_stats = {
        "Total Number of Contracts processed": total_input_files,
        "# of contracts successfully processed for AC fields": processed_ac_files,
        "# of contracts successfully processed for B fields": processed_b_files,
        "# of contracts failed for AC fields": failed_ac_files,
        "# of contracts failed for B fields": failed_b_files,
        "# of rows for AC run": len(ac_df),
        "# of rows for B run": len(b_df),
        "Total rows in merged DataFrame": len(merged_df),
        "Unique Filenames in AC DataFrame": ac_df[file_name_column].nunique(),
        "Unique Filenames in B DataFrame": b_df[file_name_column].nunique(),
        "Unique Filenames in merged DataFrame": merged_df[file_name_column].nunique(),
        "Contracts with all AC fields missing": merged_df[
            [col for col in ac_df.columns if col != file_name_column]
        ]
        .isnull()
        .all(axis=1)
        .sum(),
        "Contracts with all B fields missing": merged_df[
            [col for col in b_df.columns if col != file_name_column]
        ]
        .isnull()
        .all(axis=1)
        .sum(),
        "# of duplicate rows for AC run": ac_df.duplicated().sum(),
        "# of duplicate rows for B run": b_df.duplicated().sum(),
        "Number of duplicate rows in merged DataFrame": merged_df.duplicated().sum(),
    }

    ac_stats = {}
    for field in [col for col in ac_df.columns if col != file_name_column]:
        missing_contracts = merged_df[merged_df[field].isnull()][
            file_name_column
        ].nunique()
        ac_stats[field] = {
            "Contracts missing this field (Merged DataFrame)": missing_contracts,
            **detect_anomalies(ac_df[field], field, is_merged=False),
            **detect_anomalies(merged_df[field], field, is_merged=True),
        }

    b_stats = {}
    for field in [col for col in b_df.columns if col != file_name_column]:
        missing_contracts = (
            merged_df.groupby(file_name_column)[field]
            .apply(lambda x: x.isnull().all())
            .sum()
        )
        contracts_with_missing = (
            merged_df.groupby(file_name_column)[field]
            .apply(lambda x: x.isnull().any())
            .sum()
        )
        total_missing = merged_df[field].isnull().sum()
        b_stats[field] = {
            "Contracts missing all entries (Merged DataFrame)": missing_contracts,
            "Contracts with at least one missing entry (Merged DataFrame)": contracts_with_missing,
            "Total missing values (Merged DataFrame)": total_missing,
            **detect_anomalies(b_df[field], field, is_merged=False),
            **detect_anomalies(merged_df[field], field, is_merged=True),
        }

    return general_stats, ac_stats, b_stats


def perform_qc_qa_tests_ac(ac_df, input_dict, processed_files):
    total_input_files = len(input_dict) if input_dict else 0
    processed_ac_files = len(
        [f for f in processed_files if f in ac_df[file_name_column].values]
    )
    failed_ac_files = (
        total_input_files - processed_ac_files if total_input_files > 0 else 0
    )

    general_stats = {
        "Total Number of Contracts processed": total_input_files,
        "# of contracts successfully processed for AC fields": processed_ac_files,
        "# of contracts failed for AC fields": failed_ac_files,
        "# of rows for AC run": len(ac_df),
        "Unique Filenames in AC DataFrame": ac_df[file_name_column].nunique(),
        "Contracts with all AC fields missing": ac_df[
            [col for col in ac_df.columns if col != file_name_column]
        ]
        .isnull()
        .all(axis=1)
        .sum(),
        "# of duplicate rows for AC run": ac_df.duplicated().sum(),
    }

    ac_stats = {}
    for field in [col for col in ac_df.columns if col != file_name_column]:
        missing_contracts = ac_df[ac_df[field].isnull()][file_name_column].nunique()
        ac_stats[field] = {
            "Contracts missing this field (Merged DataFrame)": missing_contracts,
            **detect_anomalies(ac_df[field], field, is_merged=False),
        }

    return general_stats, ac_stats


def perform_qc_qa_tests_b(b_df, input_dict, processed_files):
    total_input_files = len(input_dict) if input_dict else 0
    processed_b_files = len(
        [f for f in processed_files if f in b_df[file_name_column].values]
    )
    failed_b_files = (
        total_input_files - processed_b_files if total_input_files > 0 else 0
    )

    general_stats = {
        "Total Number of Contracts processed": total_input_files,
        "# of contracts successfully processed for B fields": processed_b_files,
        "# of contracts failed for B fields": failed_b_files,
        "# of rows for B run": len(b_df),
        "Unique Filenames in B DataFrame": b_df[file_name_column].nunique(),
        "Contracts with all B fields missing": b_df[
            [col for col in b_df.columns if col != file_name_column]
        ]
        .isnull()
        .all(axis=1)
        .sum(),
        "# of duplicate rows for B run": b_df.duplicated().sum(),
    }

    b_stats = {}
    for field in [col for col in b_df.columns if col != file_name_column]:
        missing_contracts = (
            b_df.groupby(file_name_column)[field]
            .apply(lambda x: x.isnull().all())
            .sum()
        )
        contracts_with_missing = (
            b_df.groupby(file_name_column)[field]
            .apply(lambda x: x.isnull().any())
            .sum()
        )
        total_missing = b_df[field].isnull().sum()
        b_stats[field] = {
            "Contracts missing all entries (B DataFrame)": missing_contracts,
            "Contracts with at least one missing entry (B DataFrame)": contracts_with_missing,
            "Total missing values (Merged DataFrame)": total_missing,
            **detect_anomalies(b_df[field], field, is_merged=False),
        }

    return general_stats, b_stats


def split_into_pages(text):
    text_list = text.split("Start of Page No. = ")
    text_dict = {page.split()[0]: page for page in text_list[1:]}
    return text_dict


def get_table_stats(text_dict):
    table_count, num_table_pages, rate_count = 0, 0, 0
    table_pages, rate_pages, rates_more_than_10 = [], [], []
    for page_num, page_text in text_dict.items():
        if utils.contains_reimbursement(page_text, page_num):
            rate_pages.append(page_num)
            rates_on_page = sum(1 for char in page_text if char in {"$", "%"})
            rate_count += rates_on_page
            if rates_on_page >= 10:
                rates_more_than_10.append(page_num)

        table_texts = re.findall(
            r"-------Table Start--------(.*?)-------Table End--------",
            page_text,
            re.DOTALL,
        )
        table_page_count = len(table_texts)
        if table_page_count > 0:
            table_count += table_page_count
            num_table_pages += 1
            table_pages.append(page_num)

    return {
        "Table Count": table_count,
        "Table Page Count": num_table_pages,
        "Table Pages": table_pages,
        "Rate Count": rate_count,
        "Rate Pages": rate_pages,
        "Num Rate Pages": len(rate_pages),
        ">=10 Rate Pages": rates_more_than_10,
        "Num >=10 Rates": len(rates_more_than_10),
    }


def main():
    args = parse_arguments()

    script_dir = os.path.dirname(os.path.abspath(__file__))
    parent_dir = os.path.dirname(script_dir)

    ac_df_path = (
        os.path.abspath(os.path.join(parent_dir, args.ac_df)) if args.ac_df else None
    )
    b_df_path = (
        os.path.abspath(os.path.join(parent_dir, args.b_df)) if args.b_df else None
    )

    input_dir = (
        os.path.abspath(os.path.join(parent_dir, args.input_dir))
        if args.input_dir
        else None
    )
    output_dir = (
        os.path.abspath(os.path.join(parent_dir, args.output_dir))
        if args.output_dir
        else parent_dir
    )

    reports_dir = os.path.join(parent_dir, "reports")
    os.makedirs(reports_dir, exist_ok=True)

    if ac_df_path:
        print(f"AC DataFrame: {ac_df_path}")
    if b_df_path:
        print(f"B DataFrame: {b_df_path}")
    if input_dir:
        print(f"Input Dir: {input_dir}")
    print(f"Output Dir: {output_dir}")
    print(f"Reports Dir: {reports_dir}")

    if not input_dir:
        print("Needs an Input Directory!")
        return
    if not ac_df_path and not b_df_path:
        print("Needs at least one of AC Outputs and B Outputs!")
        return

    input_dict = read_input_files(input_dir)
    total_files = len(input_dict)
    print(f"Total Input Files: {total_files}")

    all_input_filenames = (
        [remove_txt_extension(f.lower()) for f in input_dict.keys()]
        if input_dict
        else []
    )

    ac_df = (
        pd.read_csv(ac_df_path).drop(columns=["Unnamed: 0"], errors="ignore")
        if ac_df_path
        else None
    )
    b_df = (
        pd.read_csv(b_df_path).drop(columns=["Unnamed: 0"], errors="ignore")
        if b_df_path
        else None
    )

    if ac_df_path:
        ac_df[file_name_column] = ac_df[file_name_column].apply(
            lambda x: remove_txt_extension(str(x)).lower()
        )
        ac_filenames = set(ac_df[file_name_column])
    if b_df_path:
        b_filenames = set(b_df[file_name_column])
        b_df[file_name_column] = b_df[file_name_column].apply(
            lambda x: remove_txt_extension(str(x)).lower()
        )

    if ac_df_path and b_df_path:
        unmatched_b_filenames = b_filenames - ac_filenames
        unmatched_b_df = pd.DataFrame({file_name_column: list(unmatched_b_filenames)})
        unmatched_b_output_path = os.path.join(output_dir, "unmatched_B_filenames.csv")
        unmatched_b_df.to_csv(unmatched_b_output_path, index=False)
        print(f"Unmatched B filenames saved to '{unmatched_b_output_path}'")

        # Perform an inner merge to keep only matching rows for the main output
        merged_df = pd.merge(ac_df, b_df, on=file_name_column, how="inner")
        merged_df = merged_df.drop(columns=["Unnamed: 0"], errors="ignore")
        ac_rows_before = len(ac_df)
        b_rows_before = len(b_df)
        merged_rows = len(merged_df)
        unmatched_b_count = len(unmatched_b_filenames)

        print(f"Rows in AC DataFrame before merge: {ac_rows_before}")
        print(f"Rows in B DataFrame before merge: {b_rows_before}")
        print(f"Rows in merged DataFrame: {merged_rows}")
        print(f"Unique filenames in B not found in AC: {unmatched_b_count}")

        merged_output_path = os.path.join(output_dir, "DOCZYAI_merged_results.csv")
        merged_df.to_csv(merged_output_path, index=False)
        print(f"Merged data saved to '{merged_output_path}'")

        processed_files = list(
            set(ac_df[file_name_column].tolist() + b_df[file_name_column].tolist())
        )
        general_stats, ac_stats, b_stats = perform_qc_qa_tests_both(
            ac_df, b_df, merged_df, input_dict, processed_files
        )

    elif ac_df_path:
        processed_files = list(set(ac_df[file_name_column].tolist()))
        general_stats, ac_stats = perform_qc_qa_tests_ac(
            ac_df, input_dict, processed_files
        )
    else:
        processed_files = list(set(b_df[file_name_column].tolist()))
        general_stats, b_stats = perform_qc_qa_tests_b(
            b_df, input_dict, processed_files
        )

    table_stats = []
    if input_dict:
        for filename, content in input_dict.items():
            filename_without_ext = remove_txt_extension(filename).lower()
            text_dict = split_into_pages(content)
            stats = get_table_stats(text_dict)
            table_stats.append(
                {
                    "Contract Name": filename,
                    "Has B output": (
                        filename_without_ext in b_df[file_name_column].values
                        if b_df_path
                        else "No B Output Given"
                    ),
                    "Has AC output": (
                        filename_without_ext in ac_df[file_name_column].values
                        if ac_df_path
                        else "No AC Output Given"
                    ),
                    "# Tables": stats["Table Count"],
                    "# Reimb rates": stats["Rate Count"],
                    "Table Pages": ", ".join(map(str, stats["Table Pages"])),
                    "Rate Pages": ", ".join(map(str, stats["Rate Pages"])),
                    ">=10 Rate Pages": ", ".join(map(str, stats[">=10 Rate Pages"])),
                }
            )

    wb = Workbook()
    ws_general = wb.active
    ws_general.title = "General Stats"

    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    general_stats_with_metadata = {
        "Report Generation Time": timestamp,
        "AC File": args.ac_df if ac_df_path else "No AC Output Given",
        "B File": args.b_df if b_df_path else "No B Output Given",
        **general_stats,
    }

    general_df = pd.DataFrame.from_dict(
        general_stats_with_metadata, orient="index", columns=["Value"]
    )
    for r in dataframe_to_rows(general_df, index=True, header=True):
        ws_general.append(r)

    if ac_df_path:
        ws_ac = wb.create_sheet("AC Fields Stats")
        ac_stats_df = pd.DataFrame(ac_stats).T
        for r in dataframe_to_rows(ac_stats_df, index=True, header=True):
            ws_ac.append(r)

    if b_df_path:
        ws_b = wb.create_sheet("B Fields Stats")
        b_stats_df = pd.DataFrame(b_stats).T
        for r in dataframe_to_rows(b_stats_df, index=True, header=True):
            ws_b.append(r)

    if table_stats:
        ws_table_stats = wb.create_sheet("Table and Reimbursement Stats")
        table_stats_df = pd.DataFrame(table_stats)
        for r in dataframe_to_rows(table_stats_df, index=False, header=True):
            ws_table_stats.append(r)

    excel_report_filename = os.path.join(
        reports_dir, f'QC_QA_Report_{datetime.now().strftime("%Y%m%d_%H%M%S")}.xlsx'
    )
    wb.save(excel_report_filename)
    print(f"Excel QC/QA Report saved to {excel_report_filename}")


if __name__ == "__main__":
    main()