doczyai-pipelines/src/scripts/missingness_analysis.py

import src.config as config
from datetime import datetime
import pandas as pd
import src.utils.string_utils as string_utils
import io


def find_inv_test_results():
    """
    Efficiently find the 30 most recent directories containing 'inv-test'
    and list their RESULTS.csv files.

    Strategy: Only list top-level directories first, filter for 'inv-test',
    then search only those directories for RESULTS.csv files.
    """
    s3_client = config.S3_CLIENT
    bucket_name = config.S3_OUTPUT_BUCKET

    print(f"\nSearching for 'inv-test' directories in bucket: {bucket_name}")
    print("=" * 80)

    # Step 1: Get only top-level directories (much faster than scanning all files)
    print("Step 1: Listing top-level directories...")
    response = s3_client.list_objects_v2(Bucket=bucket_name, Delimiter="/")

    if "CommonPrefixes" not in response:
        print("No directories found in bucket")
        return []

    all_top_level_dirs = [
        prefix["Prefix"].rstrip("/") for prefix in response["CommonPrefixes"]
    ]
    print(f"Found {len(all_top_level_dirs)} total top-level directories")

    # Step 2: Filter for 'inv-test' directories only
    inv_test_dirs = [d for d in all_top_level_dirs if "inv-test" in d]
    print(f"Found {len(inv_test_dirs)} directories containing 'inv-test'")

    if not inv_test_dirs:
        print("No 'inv-test' directories found")
        return []

    # Step 3: Get timestamps for inv-test directories only (not all files in bucket)
    print("\nStep 2: Getting timestamps for inv-test directories...")
    dir_timestamps = []

    for dir_name in inv_test_dirs:
        # Get just the first file in each directory to get a timestamp
        response = s3_client.list_objects_v2(
            Bucket=bucket_name, Prefix=f"{dir_name}/", MaxKeys=1
        )

        if "Contents" in response and len(response["Contents"]) > 0:
            timestamp = response["Contents"][0]["LastModified"]
            dir_timestamps.append((dir_name, timestamp))

    # Sort by timestamp (most recent first)
    dir_timestamps.sort(key=lambda x: x[1], reverse=True)

    # Get the 30 most recent
    recent_inv_test_dirs = dir_timestamps[:30]

    print(
        f"\nStep 3: Searching 30 most recent 'inv-test' directories for RESULTS.csv files:"
    )
    print("=" * 80)

    total_results_files = 0
    all_results_files = {}  # Dictionary to store directory -> list of file paths

    for i, (dir_name, timestamp) in enumerate(recent_inv_test_dirs, 1):
        print(f"\n[{i}/30] Directory: {dir_name}")
        print(f"        Last Modified: {timestamp.strftime('%Y-%m-%d %H:%M:%S')}")

        # Search only this directory for RESULTS.csv files
        results_files = []
        paginator = s3_client.get_paginator("list_objects_v2")

        for page in paginator.paginate(Bucket=bucket_name, Prefix=f"{dir_name}/"):
            if "Contents" not in page:
                continue

            for obj in page["Contents"]:
                key = obj["Key"]
                # Ignore files in subdirectories named "individual"
                if "/individual/" in key:
                    continue
                if key.endswith("-RESULTS.csv"):
                    results_files.append(key)

        if results_files:
            print(f"        Found {len(results_files)} RESULTS.csv file(s):")
            for file_key in results_files:
                print(f"          ✓ {file_key}")
                total_results_files += 1
            all_results_files[dir_name] = results_files
        else:
            print(f"        ✗ No -RESULTS.csv files found")

    print("\n" + "=" * 80)
    print(f"Summary:")
    print(f"  - Searched 30 most recent 'inv-test' directories")
    print(f"  - Total -RESULTS.csv files found: {total_results_files}")
    print("=" * 80)

    # Step 4: Generate missingness report
    if total_results_files > 0:
        print(f"\nStep 4: Generating missingness report...")
        generate_missingness_report(s3_client, bucket_name, all_results_files)

    return recent_inv_test_dirs


def generate_missingness_report(s3_client, bucket_name, results_files_dict):
    """
    Generate a missingness report for all RESULTS.csv files.

    Args:
        s3_client: S3 client object
        bucket_name: S3 bucket name
        results_files_dict: Dictionary mapping directory names to lists of result file paths
    """
    print("=" * 80)
    print("Generating Missingness Report")
    print("=" * 80)

    missingness_data = []

    for dir_name, file_paths in results_files_dict.items():
        for file_path in file_paths:
            print(f"\nProcessing: {file_path}")

            try:
                # Read CSV from S3
                response = s3_client.get_object(Bucket=bucket_name, Key=file_path)
                csv_content = response["Body"].read().decode("utf-8")
                df = pd.read_csv(io.StringIO(csv_content))

                if df.empty:
                    print(f"  ⚠ File is empty, skipping")
                    continue

                # Calculate missingness for each column
                row_data = {"csv_file": file_path}

                for column in df.columns:
                    empty_mask = string_utils.is_empty(df[column], pd_mask=True)
                    missingness_pct = (empty_mask.sum() / len(df)) * 100
                    row_data[column] = round(missingness_pct, 2)

                missingness_data.append(row_data)
                print(f"  ✓ Processed {len(df)} rows, {len(df.columns)} columns")

            except Exception as e:
                print(f"  ✗ Error processing file: {e}")
                continue

    if not missingness_data:
        print("\nNo valid data to create report")
        return

    # Create missingness DataFrame
    missingness_df = pd.DataFrame(missingness_data)

    # Set csv_file as index for better readability
    missingness_df.set_index("csv_file", inplace=True)

    # Save report
    output_filename = (
        f"missingness_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
    )
    missingness_df.to_csv(output_filename)

    print("\n" + "=" * 80)
    print(f"Missingness Report Summary:")
    print(f"  - Total CSV files processed: {len(missingness_data)}")
    print(f"  - Total columns analyzed: {len(missingness_df.columns)}")
    print(f"  - Report saved to: {output_filename}")
    print("=" * 80)

    # Display sample of the report
    print("\nSample of Missingness Report (first 5 columns):")
    print("-" * 80)
    if len(missingness_df.columns) > 5:
        print(missingness_df.iloc[:, :5].to_string())
        print(f"\n... and {len(missingness_df.columns) - 5} more columns")
    else:
        print(missingness_df.to_string())
    print("-" * 80)


if __name__ == "__main__":
    find_inv_test_results()