#!/usr/bin/env python3
"""
Step 3: Exploratory Data Analysis (EDA)
==========================================
This script performs exploratory data analysis on the processed district-level data,
generating descriptive statistics and visualizations to understand distributions
and relationships between infrastructure, labor force participation, and empowerment.

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

import sys
from pathlib import Path
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from typing import Tuple, Optional

# Configure matplotlib for non-interactive backend
plt.switch_backend('Agg')

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

# Configure visualization style
sns.set_style("whitegrid")
plt.rcParams['font.family'] = 'sans-serif'
plt.rcParams['font.size'] = 10
plt.rcParams['axes.linewidth'] = 0.8
plt.rcParams['figure.dpi'] = 100

# Define paths (absolute paths for sandbox environment)
BASE_DIR = Path("/app/sandbox/session_20251229_112053_41c35bf38d4f")
DATA_DIR = BASE_DIR / "data" / "processed"
FIGURES_DIR = BASE_DIR / "figures"
RESULTS_DIR = BASE_DIR / "results"

# Input file
INPUT_FILE = DATA_DIR / "analysis_ready.csv"

# Key variables for analysis
KEY_VARIABLES = ['Infrastructure_Index', 'FLFPR', 'Empowerment_Index']


def validate_input() -> bool:
    """
    Validate that the input file exists.

    Returns:
        bool: True if input exists, False otherwise
    """
    if not INPUT_FILE.exists():
        print("=" * 70)
        print("ERROR: Input file not found!")
        print("=" * 70)
        print(f"\nExpected file: {INPUT_FILE}")
        print("\nThis script requires the output from Step 2 (Data Preprocessing).")
        print("Please run the following command first:")
        print("\n  python3 workflow/02_data_preprocessing.py")
        print("\nThen re-run this script.")
        print("=" * 70)
        return False
    return True


def load_data() -> pd.DataFrame:
    """
    Load the processed dataset.

    Returns:
        pd.DataFrame: Loaded dataset
    """
    print(f"Loading data from: {INPUT_FILE}")
    df = pd.read_csv(INPUT_FILE)
    print(f"✓ Data loaded successfully")
    print(f"  Shape: {df.shape[0]} rows × {df.shape[1]} columns")
    print(f"  Columns: {list(df.columns)}")
    return df


def calculate_descriptive_stats(df: pd.DataFrame) -> pd.DataFrame:
    """
    Calculate descriptive statistics for key variables.

    Args:
        df: Input dataframe

    Returns:
        pd.DataFrame: Descriptive statistics
    """
    print("\nCalculating descriptive statistics...")

    # Check which key variables are available
    available_vars = [var for var in KEY_VARIABLES if var in df.columns]
    missing_vars = [var for var in KEY_VARIABLES if var not in df.columns]

    if missing_vars:
        print(f"⚠️  Warning: The following key variables are missing: {missing_vars}")

    if not available_vars:
        print("❌ ERROR: None of the key variables found in dataset!")
        print(f"   Expected: {KEY_VARIABLES}")
        print(f"   Found: {list(df.columns)}")
        sys.exit(1)

    # Calculate statistics
    stats = df[available_vars].describe().T
    stats['median'] = df[available_vars].median()

    # Reorder columns for clarity
    stats = stats[['count', 'mean', 'median', 'std', 'min', 'max']]

    print("✓ Descriptive statistics calculated")
    print("\nSummary:")
    print(stats.to_string())

    return stats


def plot_distributions(df: pd.DataFrame) -> None:
    """
    Create distribution plots (histograms with KDE) for key variables.

    Args:
        df: Input dataframe
    """
    print("\nCreating distribution plots...")

    available_vars = [var for var in KEY_VARIABLES if var in df.columns]
    n_vars = len(available_vars)

    if n_vars == 0:
        print("⚠️  No key variables available for plotting")
        return

    # Create subplots
    fig, axes = plt.subplots(1, n_vars, figsize=(6 * n_vars, 5))
    if n_vars == 1:
        axes = [axes]

    for idx, var in enumerate(available_vars):
        ax = axes[idx]

        # Plot histogram with KDE
        data = df[var].dropna()
        ax.hist(data, bins=30, alpha=0.6, color='steelblue',
                edgecolor='black', density=True, label='Histogram')

        # Add KDE
        from scipy import stats as scipy_stats
        kde = scipy_stats.gaussian_kde(data)
        x_range = np.linspace(data.min(), data.max(), 200)
        ax.plot(x_range, kde(x_range), 'r-', linewidth=2, label='KDE')

        # Formatting
        ax.set_xlabel(var, fontsize=11, fontweight='bold')
        ax.set_ylabel('Density', fontsize=11)
        ax.set_title(f'Distribution of {var}', fontsize=12, fontweight='bold')
        ax.legend()
        ax.grid(True, alpha=0.3)

    plt.tight_layout()
    output_file = FIGURES_DIR / "01_distributions.png"
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    plt.close()

    print(f"✓ Distribution plots saved to: {output_file}")


def plot_correlation_matrix(df: pd.DataFrame) -> None:
    """
    Create correlation heatmap for all numeric variables.

    Args:
        df: Input dataframe
    """
    print("\nCreating correlation matrix...")

    # Select numeric columns only
    numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()

    if len(numeric_cols) < 2:
        print("⚠️  Not enough numeric variables for correlation analysis")
        return

    # Calculate correlation matrix
    corr_matrix = df[numeric_cols].corr()

    # Create heatmap
    fig, ax = plt.subplots(figsize=(10, 8))

    # Use diverging colormap
    mask = np.triu(np.ones_like(corr_matrix, dtype=bool), k=1)  # Mask upper triangle
    sns.heatmap(corr_matrix, mask=mask, annot=True, fmt='.2f',
                cmap='RdBu_r', center=0, vmin=-1, vmax=1,
                square=True, linewidths=0.5, cbar_kws={"shrink": 0.8},
                ax=ax)

    ax.set_title('Correlation Matrix (Lower Triangle)',
                 fontsize=14, fontweight='bold', pad=20)

    plt.tight_layout()
    output_file = FIGURES_DIR / "02_correlation_matrix.png"
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    plt.close()

    print(f"✓ Correlation matrix saved to: {output_file}")

    # Print strong correlations
    print("\nStrong correlations (|r| > 0.5):")
    strong_corr = []
    for i in range(len(corr_matrix.columns)):
        for j in range(i+1, len(corr_matrix.columns)):
            corr_val = corr_matrix.iloc[i, j]
            if abs(corr_val) > 0.5:
                strong_corr.append((corr_matrix.columns[i],
                                  corr_matrix.columns[j],
                                  corr_val))

    if strong_corr:
        for var1, var2, corr_val in sorted(strong_corr,
                                           key=lambda x: abs(x[2]),
                                           reverse=True):
            print(f"  {var1} ↔ {var2}: r = {corr_val:.3f}")
    else:
        print("  None found")


def plot_scatter_relationships(df: pd.DataFrame) -> None:
    """
    Create scatter plots to visualize relationships between key variables.

    Args:
        df: Input dataframe
    """
    print("\nCreating scatter plots...")

    # Define pairs to plot
    pairs = [
        ('Infrastructure_Index', 'FLFPR'),
        ('Infrastructure_Index', 'Empowerment_Index'),
    ]

    # Filter to available pairs
    available_pairs = []
    for x_var, y_var in pairs:
        if x_var in df.columns and y_var in df.columns:
            available_pairs.append((x_var, y_var))

    if not available_pairs:
        print("⚠️  No variable pairs available for scatter plots")
        return

    n_pairs = len(available_pairs)

    # Check if State column exists for color-coding
    has_state = 'State' in df.columns or 'state' in df.columns
    state_col = 'State' if 'State' in df.columns else ('state' if 'state' in df.columns else None)

    # Create subplots
    fig, axes = plt.subplots(1, n_pairs, figsize=(7 * n_pairs, 6))
    if n_pairs == 1:
        axes = [axes]

    for idx, (x_var, y_var) in enumerate(available_pairs):
        ax = axes[idx]

        # Prepare data
        plot_data = df[[x_var, y_var]].dropna()

        if has_state and state_col:
            # Color by state if available
            plot_data_with_state = df[[x_var, y_var, state_col]].dropna()

            # If too many states, just use plain scatter
            n_states = plot_data_with_state[state_col].nunique()
            if n_states <= 40:  # Reasonable number for distinct colors
                for state in plot_data_with_state[state_col].unique():
                    state_data = plot_data_with_state[
                        plot_data_with_state[state_col] == state
                    ]
                    ax.scatter(state_data[x_var], state_data[y_var],
                             alpha=0.6, s=60, label=state if n_states <= 10 else None)
            else:
                ax.scatter(plot_data[x_var], plot_data[y_var],
                         alpha=0.6, s=60, color='steelblue')
        else:
            ax.scatter(plot_data[x_var], plot_data[y_var],
                     alpha=0.6, s=60, color='steelblue')

        # Add trend line
        from scipy import stats as scipy_stats
        if len(plot_data) >= 3:
            z = np.polyfit(plot_data[x_var], plot_data[y_var], 1)
            p = np.poly1d(z)
            x_line = np.linspace(plot_data[x_var].min(),
                               plot_data[x_var].max(), 100)
            ax.plot(x_line, p(x_line), "r--", alpha=0.8, linewidth=2,
                   label='Trend line')

            # Calculate correlation
            r, p_val = scipy_stats.pearsonr(plot_data[x_var], plot_data[y_var])
            ax.text(0.05, 0.95, f'r = {r:.3f}\np = {p_val:.3e}',
                   transform=ax.transAxes, fontsize=10,
                   verticalalignment='top',
                   bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))

        # Formatting
        ax.set_xlabel(x_var, fontsize=11, fontweight='bold')
        ax.set_ylabel(y_var, fontsize=11, fontweight='bold')
        ax.set_title(f'{x_var} vs {y_var}', fontsize=12, fontweight='bold')
        ax.grid(True, alpha=0.3)

        # Add legend only if states are plotted and there aren't too many
        if has_state and state_col:
            n_states = df[state_col].nunique()
            if n_states <= 10:
                ax.legend(fontsize=8, loc='best')

    plt.tight_layout()
    output_file = FIGURES_DIR / "03_scatter_relationships.png"
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    plt.close()

    print(f"✓ Scatter plots saved to: {output_file}")


def save_results(stats: pd.DataFrame) -> None:
    """
    Save descriptive statistics to CSV.

    Args:
        stats: Descriptive statistics dataframe
    """
    print("\nSaving results...")
    output_file = RESULTS_DIR / "descriptive_stats.csv"
    stats.to_csv(output_file)
    print(f"✓ Descriptive statistics saved to: {output_file}")


def main():
    """Main execution function."""
    print("=" * 70)
    print("STEP 3: EXPLORATORY DATA ANALYSIS (EDA)")
    print("=" * 70)
    print()

    # Step 1: Validate input
    if not validate_input():
        sys.exit(1)

    print()

    # Step 2: Load data
    df = load_data()

    # Step 3: Calculate and save descriptive statistics
    stats = calculate_descriptive_stats(df)
    save_results(stats)

    # Step 4: Create visualizations
    print("\n" + "=" * 70)
    print("GENERATING VISUALIZATIONS")
    print("=" * 70)

    plot_distributions(df)
    plot_correlation_matrix(df)
    plot_scatter_relationships(df)

    # Summary
    print("\n" + "=" * 70)
    print("EDA COMPLETE")
    print("=" * 70)
    print(f"\nOutputs:")
    print(f"  📊 Statistics: {RESULTS_DIR / 'descriptive_stats.csv'}")
    print(f"  📈 Distributions: {FIGURES_DIR / '01_distributions.png'}")
    print(f"  🔥 Correlation Matrix: {FIGURES_DIR / '02_correlation_matrix.png'}")
    print(f"  📉 Scatter Plots: {FIGURES_DIR / '03_scatter_relationships.png'}")
    print()
    print("Next step: Run statistical analysis (Step 4)")
    print("=" * 70)


if __name__ == "__main__":
    main()