#!/usr/bin/env python3
"""
Statistical Modeling: Infrastructure Impact on Gender Outcomes
=============================================================

This script performs OLS regression analysis to quantify the relationship
between infrastructure development and gender outcome indicators using
Census 2011 and NFHS-5 data.

Author: K-Dense Coding Agent
Date: 2025-12-29
"""

import sys
import warnings
from pathlib import Path
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')  # Non-interactive backend
import matplotlib.pyplot as plt
import seaborn as sns
from scipy import stats

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

# Suppress warnings for cleaner output
warnings.filterwarnings('ignore', category=FutureWarning)
warnings.filterwarnings('ignore', category=UserWarning)

# Configure matplotlib for publication-quality figures
plt.rcParams['font.family'] = 'sans-serif'
plt.rcParams['font.size'] = 10
plt.rcParams['axes.linewidth'] = 0.8
plt.rcParams['figure.dpi'] = 300

# Define paths
SESSION_DIR = Path('/app/sandbox/session_20251229_112053_41c35bf38d4f')
INPUT_FILE = SESSION_DIR / 'data' / 'processed' / 'analysis_ready.csv'
RESULTS_DIR = SESSION_DIR / 'results'
FIGURES_DIR = SESSION_DIR / 'figures'

# Ensure output directories exist
RESULTS_DIR.mkdir(parents=True, exist_ok=True)
FIGURES_DIR.mkdir(parents=True, exist_ok=True)


def validate_input():
    """
    Validate that input data file exists.

    Returns:
        bool: True if input file exists, False otherwise
    """
    print("=" * 70)
    print("STEP 1: INPUT VALIDATION")
    print("=" * 70)

    if not INPUT_FILE.exists():
        print(f"❌ ERROR: Input file not found: {INPUT_FILE}")
        print("\nThe analysis-ready dataset does not exist yet.")
        print("Please ensure the following steps are completed:")
        print("  1. Upload raw data files (census_2011_district.csv, nfhs_5_district.csv)")
        print("  2. Run workflow/02_data_preprocessing.py to generate processed data")
        print("\nExiting...")
        return False

    print(f"✓ Input file found: {INPUT_FILE}")
    file_size = INPUT_FILE.stat().st_size / 1024  # KB
    print(f"  File size: {file_size:.1f} KB")
    return True


def load_and_explore_data():
    """
    Load the processed dataset and perform initial exploration.

    Returns:
        pd.DataFrame: Loaded dataset
    """
    print("\n" + "=" * 70)
    print("STEP 2: DATA LOADING")
    print("=" * 70)

    df = pd.read_csv(INPUT_FILE)
    print(f"✓ Dataset loaded successfully")
    print(f"  Shape: {df.shape[0]} rows × {df.shape[1]} columns")
    print(f"\nColumn names:\n  {', '.join(df.columns.tolist())}")

    # Check for missing values
    missing = df.isnull().sum()
    if missing.sum() > 0:
        print(f"\n⚠ Missing values detected:")
        for col, count in missing[missing > 0].items():
            pct = 100 * count / len(df)
            print(f"  {col}: {count} ({pct:.1f}%)")
    else:
        print(f"\n✓ No missing values in dataset")

    return df


def identify_variables(df):
    """
    Identify and validate required variables for regression models.

    Args:
        df: Input dataframe

    Returns:
        tuple: (dependent_vars, independent_var, control_vars)
    """
    print("\n" + "=" * 70)
    print("STEP 3: VARIABLE IDENTIFICATION")
    print("=" * 70)

    # Define required dependent variables
    dependent_vars = ['FLFPR', 'Empowerment_Index']

    # Define independent variable
    independent_var = 'Infrastructure_Index'

    # Validate dependent variables
    print(f"\nDependent Variables (Outcomes):")
    missing_deps = []
    for var in dependent_vars:
        if var in df.columns:
            print(f"  ✓ {var} - Found")
        else:
            print(f"  ❌ {var} - NOT FOUND")
            missing_deps.append(var)

    # Validate independent variable
    print(f"\nIndependent Variable (Predictor):")
    if independent_var in df.columns:
        print(f"  ✓ {independent_var} - Found")
    else:
        print(f"  ❌ {independent_var} - NOT FOUND")
        sys.exit(1)

    # Search for control variables
    print(f"\nControl Variables:")
    control_candidates = {
        'Literacy': ['Literacy_Rate', 'Literacy', 'literacy_rate', 'literacy'],
        'Urbanization': ['Urban_Population_Pct', 'Urbanization', 'urban_pct', 'urban_population']
    }

    control_vars = []
    for category, candidates in control_candidates.items():
        found = False
        for candidate in candidates:
            if candidate in df.columns:
                control_vars.append(candidate)
                print(f"  ✓ {category}: {candidate} - Found")
                found = True
                break
        if not found:
            print(f"  ⚠ {category}: Not found (will proceed with unadjusted model)")

    if not control_vars:
        print(f"\n⚠ WARNING: No control variables found.")
        print(f"  Models will be unadjusted (simple regression without controls).")
    else:
        print(f"\n✓ Found {len(control_vars)} control variable(s): {', '.join(control_vars)}")

    if missing_deps:
        print(f"\n❌ ERROR: Missing required dependent variables: {', '.join(missing_deps)}")
        sys.exit(1)

    return dependent_vars, independent_var, control_vars


def prepare_regression_data(df, outcome, predictor, controls):
    """
    Prepare data for regression by handling missing values.

    Args:
        df: Input dataframe
        outcome: Dependent variable name
        predictor: Independent variable name
        controls: List of control variable names

    Returns:
        pd.DataFrame: Clean dataset for regression
    """
    # Select relevant columns
    cols = [outcome, predictor] + controls
    reg_data = df[cols].copy()

    # Remove rows with missing values
    initial_n = len(reg_data)
    reg_data = reg_data.dropna()
    final_n = len(reg_data)

    if final_n < initial_n:
        dropped = initial_n - final_n
        print(f"  Dropped {dropped} rows ({100*dropped/initial_n:.1f}%) with missing values")

    return reg_data


def run_ols_regression(df, outcome, predictor, controls):
    """
    Run OLS regression model using statsmodels.

    Args:
        df: Input dataframe
        outcome: Dependent variable name
        predictor: Independent variable name
        controls: List of control variable names

    Returns:
        statsmodels regression results object
    """
    import statsmodels.api as sm
    from statsmodels.formula.api import ols

    # Prepare data
    reg_data = prepare_regression_data(df, outcome, predictor, controls)
    print(f"  Sample size: N = {len(reg_data)}")

    # Build formula
    if controls:
        formula = f"{outcome} ~ {predictor} + " + " + ".join(controls)
    else:
        formula = f"{outcome} ~ {predictor}"

    print(f"  Formula: {formula}")

    # Fit model
    model = ols(formula, data=reg_data).fit()

    return model, reg_data


def extract_model_metrics(model, predictor):
    """
    Extract key metrics from regression model.

    Args:
        model: Fitted statsmodels regression object
        predictor: Independent variable name

    Returns:
        dict: Dictionary with key model metrics
    """
    # Get coefficient for predictor
    coef = model.params[predictor]
    pval = model.pvalues[predictor]
    ci = model.conf_int().loc[predictor]
    se = model.bse[predictor]

    metrics = {
        'Coefficient': coef,
        'Std_Error': se,
        'P_value': pval,
        'CI_Lower': ci[0],
        'CI_Upper': ci[1],
        'R_squared': model.rsquared,
        'Adj_R_squared': model.rsquared_adj,
        'F_statistic': model.fvalue,
        'F_pvalue': model.f_pvalue,
        'N_obs': int(model.nobs)
    }

    return metrics


def save_regression_summaries(models_dict, summary_file):
    """
    Save full regression summaries to text file.

    Args:
        models_dict: Dictionary mapping outcome names to fitted models
        summary_file: Path to output file
    """
    with open(summary_file, 'w') as f:
        f.write("=" * 80 + "\n")
        f.write("REGRESSION ANALYSIS SUMMARIES\n")
        f.write("Infrastructure Impact on Gender Outcomes\n")
        f.write("=" * 80 + "\n\n")

        for outcome, (model, reg_data) in models_dict.items():
            f.write("=" * 80 + "\n")
            f.write(f"Model: {outcome}\n")
            f.write("=" * 80 + "\n\n")
            f.write(str(model.summary()))
            f.write("\n\n")

    print(f"✓ Saved full regression summaries to: {summary_file}")


def save_model_results_table(metrics_df, results_file):
    """
    Save key model results to CSV.

    Args:
        metrics_df: DataFrame with model metrics
        results_file: Path to output CSV file
    """
    metrics_df.to_csv(results_file, index=False, float_format='%.6f')
    print(f"✓ Saved model results table to: {results_file}")


def plot_coefficient_forest(metrics_df, output_file):
    """
    Generate forest plot showing regression coefficients with confidence intervals.

    Args:
        metrics_df: DataFrame with model metrics
        output_file: Path to save figure
    """
    fig, ax = plt.subplots(figsize=(8, 6))

    # Prepare data for plotting
    outcomes = metrics_df['Outcome'].values
    coefs = metrics_df['Coefficient'].values
    ci_lower = metrics_df['CI_Lower'].values
    ci_upper = metrics_df['CI_Upper'].values
    pvals = metrics_df['P_value'].values

    # Calculate error bars (distance from coefficient to CI bounds)
    lower_err = coefs - ci_lower
    upper_err = ci_upper - coefs

    # Create color based on significance
    colors = ['#2ecc71' if p < 0.05 else '#95a5a6' for p in pvals]

    # Plot
    y_positions = np.arange(len(outcomes))
    ax.errorbar(coefs, y_positions, xerr=[lower_err, upper_err],
                fmt='o', markersize=8, capsize=5, capthick=2,
                color='black', ecolor=colors, elinewidth=2)

    # Vertical line at zero
    ax.axvline(x=0, color='red', linestyle='--', linewidth=1, alpha=0.7)

    # Formatting
    ax.set_yticks(y_positions)
    ax.set_yticklabels(outcomes)
    ax.set_xlabel('Coefficient (Effect Size)', fontsize=12, fontweight='bold')
    ax.set_ylabel('Outcome Variable', fontsize=12, fontweight='bold')
    ax.set_title('Infrastructure Impact on Gender Outcomes\nRegression Coefficients with 95% CI',
                 fontsize=13, fontweight='bold', pad=15)
    ax.grid(axis='x', alpha=0.3, linestyle='--')

    # Add significance annotation
    for i, (coef, pval) in enumerate(zip(coefs, pvals)):
        sig_marker = '***' if pval < 0.001 else '**' if pval < 0.01 else '*' if pval < 0.05 else 'ns'
        ax.text(coef, i + 0.15, sig_marker, ha='center', fontsize=10, fontweight='bold')

    # Legend
    from matplotlib.lines import Line2D
    legend_elements = [
        Line2D([0], [0], marker='o', color='w', markerfacecolor='black',
               markersize=8, label='Coefficient'),
        Line2D([0], [0], color='#2ecc71', linewidth=2, label='95% CI (p < 0.05)'),
        Line2D([0], [0], color='#95a5a6', linewidth=2, label='95% CI (p ≥ 0.05)')
    ]
    ax.legend(handles=legend_elements, loc='best', frameon=True, fontsize=9)

    plt.tight_layout()
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"✓ Saved coefficient plot to: {output_file}")


def plot_diagnostic_residuals(models_dict, output_file):
    """
    Generate diagnostic plots (Residuals vs Fitted) for model assumptions.

    Args:
        models_dict: Dictionary mapping outcome names to fitted models
        output_file: Path to save figure
    """
    n_models = len(models_dict)
    fig, axes = plt.subplots(n_models, 2, figsize=(12, 4 * n_models))

    if n_models == 1:
        axes = axes.reshape(1, -1)

    for i, (outcome, (model, reg_data)) in enumerate(models_dict.items()):
        # Get residuals and fitted values
        residuals = model.resid
        fitted = model.fittedvalues

        # Standardized residuals
        std_resid = residuals / np.std(residuals)

        # Plot 1: Residuals vs Fitted
        axes[i, 0].scatter(fitted, residuals, alpha=0.6, s=30, edgecolors='black', linewidths=0.5)
        axes[i, 0].axhline(y=0, color='red', linestyle='--', linewidth=2)
        axes[i, 0].set_xlabel('Fitted Values', fontsize=10, fontweight='bold')
        axes[i, 0].set_ylabel('Residuals', fontsize=10, fontweight='bold')
        axes[i, 0].set_title(f'{outcome}: Residuals vs Fitted', fontsize=11, fontweight='bold')
        axes[i, 0].grid(alpha=0.3, linestyle='--')

        # Plot 2: Q-Q Plot (Normal Probability Plot)
        stats.probplot(residuals, dist="norm", plot=axes[i, 1])
        axes[i, 1].set_title(f'{outcome}: Normal Q-Q Plot', fontsize=11, fontweight='bold')
        axes[i, 1].grid(alpha=0.3, linestyle='--')

    plt.tight_layout()
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"✓ Saved diagnostic plots to: {output_file}")


def main():
    """Main execution function."""
    print("\n")
    print("#" * 70)
    print("# STATISTICAL MODELING: INFRASTRUCTURE IMPACT ON GENDER OUTCOMES")
    print("#" * 70)
    print("\n")

    # Step 1: Validate input
    if not validate_input():
        return

    # Step 2: Load data
    df = load_and_explore_data()

    # Step 3: Identify variables
    dependent_vars, independent_var, control_vars = identify_variables(df)

    # Step 4: Run regression models
    print("\n" + "=" * 70)
    print("STEP 4: REGRESSION ANALYSIS")
    print("=" * 70)

    models_dict = {}
    all_metrics = []

    for outcome in dependent_vars:
        print(f"\n--- Model {len(models_dict) + 1}: {outcome} ---")
        try:
            model, reg_data = run_ols_regression(df, outcome, independent_var, control_vars)
            models_dict[outcome] = (model, reg_data)

            # Extract metrics
            metrics = extract_model_metrics(model, independent_var)
            metrics['Outcome'] = outcome
            metrics['Model'] = f"{outcome} ~ Infrastructure_Index" + (" + Controls" if control_vars else "")
            all_metrics.append(metrics)

            # Print summary
            coef = metrics['Coefficient']
            pval = metrics['P_value']
            r2 = metrics['R_squared']
            print(f"  ✓ Model fitted successfully")
            print(f"    Coefficient: {coef:.4f} (p = {pval:.4f})")
            print(f"    R²: {r2:.4f}")

        except Exception as e:
            print(f"  ❌ Error fitting model for {outcome}: {str(e)}")
            continue

    if not models_dict:
        print("\n❌ ERROR: No models were successfully fitted.")
        sys.exit(1)

    # Step 5: Save outputs
    print("\n" + "=" * 70)
    print("STEP 5: SAVING OUTPUTS")
    print("=" * 70 + "\n")

    # Save full summaries
    summary_file = RESULTS_DIR / 'regression_summary.txt'
    save_regression_summaries(models_dict, summary_file)

    # Save metrics table
    metrics_df = pd.DataFrame(all_metrics)
    # Reorder columns for better readability
    col_order = ['Outcome', 'Model', 'Coefficient', 'Std_Error', 'P_value',
                 'CI_Lower', 'CI_Upper', 'R_squared', 'Adj_R_squared',
                 'F_statistic', 'F_pvalue', 'N_obs']
    metrics_df = metrics_df[col_order]

    results_file = RESULTS_DIR / 'model_results.csv'
    save_model_results_table(metrics_df, results_file)

    # Step 6: Generate visualizations
    print("\n" + "=" * 70)
    print("STEP 6: GENERATING VISUALIZATIONS")
    print("=" * 70 + "\n")

    # Coefficient forest plot
    coef_plot_file = FIGURES_DIR / '04_coefficient_plot.png'
    plot_coefficient_forest(metrics_df, coef_plot_file)

    # Diagnostic residuals plot
    diag_plot_file = FIGURES_DIR / '05_residuals.png'
    plot_diagnostic_residuals(models_dict, diag_plot_file)

    # Final summary
    print("\n" + "=" * 70)
    print("ANALYSIS COMPLETE")
    print("=" * 70)
    print(f"\n✓ Successfully fitted {len(models_dict)} regression model(s)")
    print(f"\nOutputs generated:")
    print(f"  • {summary_file}")
    print(f"  • {results_file}")
    print(f"  • {coef_plot_file}")
    print(f"  • {diag_plot_file}")
    print("\n")


if __name__ == "__main__":
    main()