#!/usr/bin/env python3
"""
Data Preprocessing Pipeline for Census and NFHS-5 District-Level Data

This script:
1. Validates input data availability
2. Loads and cleans Census 2011 and NFHS-5 district-level datasets
3. Standardizes district names for merging
4. Engineers features (Infrastructure Index, FLFPR, Empowerment Index)
5. Merges datasets and outputs analysis-ready data

Author: K-Dense System
Date: 2025-12-29
"""

import sys
import re
from pathlib import Path
import pandas as pd
import numpy as np

# Set random seed for reproducibility
np.random.seed(42)

# Define paths
BASE_DIR = Path("/app/sandbox/session_20251229_112053_41c35bf38d4f")
USER_DATA_DIR = BASE_DIR / "user_data"
RAW_DATA_DIR = BASE_DIR / "data" / "raw"
PROCESSED_DATA_DIR = BASE_DIR / "data" / "processed"

CENSUS_FILENAME = "census_2011_district.csv"
NFHS_FILENAME = "nfhs_5_district.csv"

OUTPUT_FILE = PROCESSED_DATA_DIR / "analysis_ready.csv"


def standardize_district_name(name):
    """
    Standardize district names for consistent merging.

    Converts to lowercase, strips whitespace, and removes non-alphanumeric characters.

    Parameters
    ----------
    name : str
        Original district name

    Returns
    -------
    str
        Standardized district name
    """
    if pd.isna(name):
        return ""

    # Convert to lowercase
    name = str(name).lower()

    # Strip leading/trailing whitespace
    name = name.strip()

    # Remove non-alphanumeric characters (keep only letters, numbers, spaces)
    name = re.sub(r'[^a-z0-9\s]', '', name)

    # Replace multiple spaces with single space
    name = re.sub(r'\s+', ' ', name)

    return name


def find_data_file(filename):
    """
    Search for data file in user_data/ and data/raw/ directories.

    Parameters
    ----------
    filename : str
        Name of the file to search for

    Returns
    -------
    Path or None
        Full path to the file if found, None otherwise
    """
    # Check user_data first
    user_path = USER_DATA_DIR / filename
    if user_path.exists():
        print(f"✓ Found {filename} in user_data/")
        return user_path

    # Check data/raw
    raw_path = RAW_DATA_DIR / filename
    if raw_path.exists():
        print(f"✓ Found {filename} in data/raw/")
        return raw_path

    return None


def load_csv_with_encoding(filepath):
    """
    Load CSV file, trying multiple encodings if needed.

    Parameters
    ----------
    filepath : Path
        Path to the CSV file

    Returns
    -------
    pd.DataFrame
        Loaded DataFrame

    Raises
    ------
    Exception
        If file cannot be loaded with any encoding
    """
    encodings = ['utf-8', 'latin-1', 'iso-8859-1', 'cp1252']

    for encoding in encodings:
        try:
            df = pd.read_csv(filepath, encoding=encoding)
            print(f"  Successfully loaded with {encoding} encoding")
            return df
        except UnicodeDecodeError:
            continue
        except Exception as e:
            print(f"  Error with {encoding} encoding: {e}")
            continue

    raise Exception(f"Could not load {filepath} with any standard encoding")


def validate_and_load_data():
    """
    Validate data availability and load datasets.

    Returns
    -------
    tuple of (pd.DataFrame, pd.DataFrame) or (None, None)
        Census and NFHS DataFrames if found, (None, None) otherwise
    """
    print("=" * 70)
    print("Step 1: Input Validation")
    print("=" * 70)

    census_path = find_data_file(CENSUS_FILENAME)
    nfhs_path = find_data_file(NFHS_FILENAME)

    if census_path is None or nfhs_path is None:
        print("\n⚠ WAITING FOR DATA UPLOAD")
        print("-" * 70)
        print("Required data files are not yet available.")
        print(f"Expected files:")
        print(f"  - {CENSUS_FILENAME}")
        print(f"  - {NFHS_FILENAME}")
        print(f"\nPlease place these files in either:")
        print(f"  - {USER_DATA_DIR}")
        print(f"  - {RAW_DATA_DIR}")
        print(f"\nRefer to MISSING_DATA.md for download instructions.")
        print("-" * 70)
        return None, None

    print("\n✓ All required data files found!")

    # Load datasets
    print("\n" + "=" * 70)
    print("Step 2: Data Loading")
    print("=" * 70)

    print(f"\nLoading Census 2011 data from {census_path.name}...")
    census_df = load_csv_with_encoding(census_path)
    print(f"  Shape: {census_df.shape}")
    print(f"  Columns: {list(census_df.columns)[:5]}..." if len(census_df.columns) > 5 else f"  Columns: {list(census_df.columns)}")

    print(f"\nLoading NFHS-5 data from {nfhs_path.name}...")
    nfhs_df = load_csv_with_encoding(nfhs_path)
    print(f"  Shape: {nfhs_df.shape}")
    print(f"  Columns: {list(nfhs_df.columns)[:5]}..." if len(nfhs_df.columns) > 5 else f"  Columns: {list(nfhs_df.columns)}")

    return census_df, nfhs_df


def preprocess_census_data(df):
    """
    Preprocess Census 2011 data.

    - Standardize district names
    - Create Infrastructure_Index (average of water, sanitation, electricity access)

    Parameters
    ----------
    df : pd.DataFrame
        Raw Census DataFrame

    Returns
    -------
    pd.DataFrame
        Preprocessed Census DataFrame
    """
    print("\n" + "=" * 70)
    print("Step 3: Census Data Preprocessing")
    print("=" * 70)

    df = df.copy()

    # Identify district name column (flexible matching)
    district_col = None
    for col in df.columns:
        if 'district' in col.lower() and 'name' in col.lower():
            district_col = col
            break

    if district_col is None:
        # Try just 'district'
        for col in df.columns:
            if col.lower() == 'district':
                district_col = col
                break

    if district_col is None:
        raise ValueError("Could not identify district name column in Census data. "
                        f"Available columns: {list(df.columns)}")

    print(f"\n✓ Identified district column: '{district_col}'")

    # Standardize district names
    print(f"  Standardizing district names...")
    df['district_name_std'] = df[district_col].apply(standardize_district_name)
    print(f"  Example: '{df[district_col].iloc[0]}' → '{df['district_name_std'].iloc[0]}'")

    # Feature Engineering: Infrastructure Index
    print(f"\n  Creating Infrastructure_Index...")

    # Try to identify infrastructure columns (flexible matching)
    water_col = None
    sanitation_col = None
    electricity_col = None

    for col in df.columns:
        col_lower = col.lower()
        if 'water' in col_lower and water_col is None:
            water_col = col
        if 'sanitation' in col_lower or 'toilet' in col_lower and sanitation_col is None:
            sanitation_col = col
        if 'electricity' in col_lower or 'electric' in col_lower and electricity_col is None:
            electricity_col = col

    # Create Infrastructure Index if columns found
    infrastructure_cols = []
    if water_col:
        infrastructure_cols.append(water_col)
        print(f"    - Water access: {water_col}")
    if sanitation_col:
        infrastructure_cols.append(sanitation_col)
        print(f"    - Sanitation: {sanitation_col}")
    if electricity_col:
        infrastructure_cols.append(electricity_col)
        print(f"    - Electricity: {electricity_col}")

    if len(infrastructure_cols) > 0:
        # Convert to numeric, handling any non-numeric values
        for col in infrastructure_cols:
            df[col] = pd.to_numeric(df[col], errors='coerce')

        df['Infrastructure_Index'] = df[infrastructure_cols].mean(axis=1)
        print(f"  ✓ Infrastructure_Index created (mean of {len(infrastructure_cols)} indicators)")
        print(f"    Range: {df['Infrastructure_Index'].min():.2f} - {df['Infrastructure_Index'].max():.2f}")
    else:
        print(f"  ⚠ Warning: Could not identify infrastructure columns")
        print(f"    Available columns: {list(df.columns)[:10]}...")
        df['Infrastructure_Index'] = np.nan

    # Select relevant columns for merging
    keep_cols = ['district_name_std', 'Infrastructure_Index']
    # Add original district name for reference
    keep_cols.insert(1, district_col)

    # Add any other numeric columns that might be useful
    numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
    for col in numeric_cols:
        if col not in keep_cols and col != 'district_name_std':
            keep_cols.append(col)

    df_clean = df[keep_cols].copy()
    print(f"\n✓ Census preprocessing complete. Shape: {df_clean.shape}")

    return df_clean


def preprocess_nfhs_data(df):
    """
    Preprocess NFHS-5 data.

    - Standardize district names
    - Extract FLFPR (Female Labor Force Participation Rate)
    - Extract Empowerment_Index

    Parameters
    ----------
    df : pd.DataFrame
        Raw NFHS DataFrame

    Returns
    -------
    pd.DataFrame
        Preprocessed NFHS DataFrame
    """
    print("\n" + "=" * 70)
    print("Step 4: NFHS-5 Data Preprocessing")
    print("=" * 70)

    df = df.copy()

    # Identify district name column
    district_col = None
    for col in df.columns:
        if 'district' in col.lower() and 'name' in col.lower():
            district_col = col
            break

    if district_col is None:
        for col in df.columns:
            if col.lower() == 'district':
                district_col = col
                break

    if district_col is None:
        raise ValueError("Could not identify district name column in NFHS data. "
                        f"Available columns: {list(df.columns)}")

    print(f"\n✓ Identified district column: '{district_col}'")

    # Standardize district names
    print(f"  Standardizing district names...")
    df['district_name_std'] = df[district_col].apply(standardize_district_name)
    print(f"  Example: '{df[district_col].iloc[0]}' → '{df['district_name_std'].iloc[0]}'")

    # Feature Engineering: Extract FLFPR and Empowerment_Index
    print(f"\n  Extracting key indicators...")

    # Try to identify FLFPR column
    flfpr_col = None
    for col in df.columns:
        col_lower = col.lower()
        if 'flfpr' in col_lower or ('female' in col_lower and 'labor' in col_lower):
            flfpr_col = col
            break

    if flfpr_col:
        df['FLFPR'] = pd.to_numeric(df[flfpr_col], errors='coerce')
        print(f"    ✓ FLFPR extracted from: {flfpr_col}")
        print(f"      Range: {df['FLFPR'].min():.2f} - {df['FLFPR'].max():.2f}")
    else:
        print(f"    ⚠ Warning: Could not identify FLFPR column")
        df['FLFPR'] = np.nan

    # Try to identify Empowerment Index
    empowerment_col = None
    for col in df.columns:
        col_lower = col.lower()
        if 'empowerment' in col_lower and 'index' in col_lower:
            empowerment_col = col
            break

    if empowerment_col:
        df['Empowerment_Index'] = pd.to_numeric(df[empowerment_col], errors='coerce')
        print(f"    ✓ Empowerment_Index extracted from: {empowerment_col}")
        print(f"      Range: {df['Empowerment_Index'].min():.2f} - {df['Empowerment_Index'].max():.2f}")
    else:
        print(f"    ⚠ Warning: Could not identify Empowerment_Index column")
        # Try to create a composite empowerment index from related columns
        empowerment_keywords = ['women', 'female', 'empowerment', 'autonomy', 'decision']
        empowerment_cols = []
        for col in df.columns:
            col_lower = col.lower()
            if any(keyword in col_lower for keyword in empowerment_keywords):
                if df[col].dtype in [np.float64, np.int64] or pd.api.types.is_numeric_dtype(df[col]):
                    empowerment_cols.append(col)

        if len(empowerment_cols) > 0:
            print(f"    Creating composite Empowerment_Index from {len(empowerment_cols)} related columns")
            for col in empowerment_cols:
                df[col] = pd.to_numeric(df[col], errors='coerce')
            df['Empowerment_Index'] = df[empowerment_cols].mean(axis=1)
        else:
            df['Empowerment_Index'] = np.nan

    # Select relevant columns for merging
    keep_cols = ['district_name_std', 'FLFPR', 'Empowerment_Index']
    keep_cols.insert(1, district_col)

    # Add any other numeric columns that might be useful
    numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
    for col in numeric_cols:
        if col not in keep_cols and col != 'district_name_std':
            keep_cols.append(col)

    df_clean = df[keep_cols].copy()
    print(f"\n✓ NFHS preprocessing complete. Shape: {df_clean.shape}")

    return df_clean


def merge_datasets(census_df, nfhs_df):
    """
    Merge Census and NFHS datasets on standardized district names.

    Uses inner join to ensure only matched districts are retained.

    Parameters
    ----------
    census_df : pd.DataFrame
        Preprocessed Census data
    nfhs_df : pd.DataFrame
        Preprocessed NFHS data

    Returns
    -------
    pd.DataFrame
        Merged analysis-ready dataset
    """
    print("\n" + "=" * 70)
    print("Step 5: Merging Datasets")
    print("=" * 70)

    print(f"\nBefore merge:")
    print(f"  Census districts: {len(census_df)}")
    print(f"  NFHS districts: {len(nfhs_df)}")

    # Check for duplicates in standardized names
    census_dups = census_df['district_name_std'].duplicated().sum()
    nfhs_dups = nfhs_df['district_name_std'].duplicated().sum()

    if census_dups > 0:
        print(f"  ⚠ Warning: {census_dups} duplicate district names in Census data")
        # Keep first occurrence
        census_df = census_df.drop_duplicates(subset='district_name_std', keep='first')

    if nfhs_dups > 0:
        print(f"  ⚠ Warning: {nfhs_dups} duplicate district names in NFHS data")
        # Keep first occurrence
        nfhs_df = nfhs_df.drop_duplicates(subset='district_name_std', keep='first')

    # Perform inner join
    merged_df = pd.merge(
        census_df,
        nfhs_df,
        on='district_name_std',
        how='inner',
        suffixes=('_census', '_nfhs')
    )

    print(f"\nAfter inner join:")
    print(f"  Merged districts: {len(merged_df)}")
    print(f"  Match rate: {100 * len(merged_df) / min(len(census_df), len(nfhs_df)):.1f}%")

    # Report unmatched districts
    unmatched_census = set(census_df['district_name_std']) - set(merged_df['district_name_std'])
    unmatched_nfhs = set(nfhs_df['district_name_std']) - set(merged_df['district_name_std'])

    if len(unmatched_census) > 0:
        print(f"\n  Unmatched Census districts: {len(unmatched_census)}")
        if len(unmatched_census) <= 5:
            print(f"    {list(unmatched_census)}")

    if len(unmatched_nfhs) > 0:
        print(f"  Unmatched NFHS districts: {len(unmatched_nfhs)}")
        if len(unmatched_nfhs) <= 5:
            print(f"    {list(unmatched_nfhs)}")

    print(f"\n✓ Merge complete. Final shape: {merged_df.shape}")
    print(f"  Columns: {list(merged_df.columns)}")

    return merged_df


def save_output(df, output_path):
    """
    Save the analysis-ready dataset to CSV.

    Parameters
    ----------
    df : pd.DataFrame
        Merged and preprocessed dataset
    output_path : Path
        Output file path
    """
    print("\n" + "=" * 70)
    print("Step 6: Saving Output")
    print("=" * 70)

    # Ensure output directory exists
    output_path.parent.mkdir(parents=True, exist_ok=True)

    # Save to CSV
    df.to_csv(output_path, index=False)

    print(f"\n✓ Analysis-ready dataset saved to:")
    print(f"  {output_path}")
    print(f"\nDataset summary:")
    print(f"  Rows: {len(df)}")
    print(f"  Columns: {len(df.columns)}")
    print(f"  Memory usage: {df.memory_usage(deep=True).sum() / 1024**2:.2f} MB")

    # Display sample statistics
    print(f"\nKey indicators (first 5 rows):")
    print(df.head())

    print(f"\nMissing values:")
    missing = df.isnull().sum()
    missing = missing[missing > 0]
    if len(missing) > 0:
        for col, count in missing.items():
            print(f"  {col}: {count} ({100*count/len(df):.1f}%)")
    else:
        print("  None - all columns complete!")


def main():
    """Main execution function."""
    print("\n" + "=" * 70)
    print("DATA PREPROCESSING PIPELINE")
    print("Census 2011 + NFHS-5 District-Level Analysis")
    print("=" * 70)

    # Step 1-2: Validate and load data
    census_df, nfhs_df = validate_and_load_data()

    if census_df is None or nfhs_df is None:
        print("\n✓ Script completed. Waiting for data upload.")
        sys.exit(0)

    # Step 3: Preprocess Census data
    census_clean = preprocess_census_data(census_df)

    # Step 4: Preprocess NFHS data
    nfhs_clean = preprocess_nfhs_data(nfhs_df)

    # Step 5: Merge datasets
    merged_df = merge_datasets(census_clean, nfhs_clean)

    # Step 6: Save output
    save_output(merged_df, OUTPUT_FILE)

    print("\n" + "=" * 70)
    print("✓ DATA PREPROCESSING COMPLETE")
    print("=" * 70)
    print(f"\nNext steps:")
    print(f"  1. Review the output file: {OUTPUT_FILE}")
    print(f"  2. Proceed to data analysis and modeling")
    print("=" * 70 + "\n")


if __name__ == "__main__":
    main()