"""
Statistical Analysis Script

Performs comprehensive statistical analysis:
1. Exploratory Data Analysis (EDA)
2. Trend Analysis
3. Correlation Analysis
4. Regression Analysis

Identifies key drivers of emigration from Indian states.
"""

import pandas as pd
import numpy as np
from pathlib import Path
from datetime import datetime
from scipy import stats
from scipy.stats import pearsonr, spearmanr
import statsmodels.api as sm
from statsmodels.formula.api import ols
import json

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

print("=" * 70)
print("STATISTICAL ANALYSIS")
print("=" * 70)
print(f"Start time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
print()

# Define paths
DATA_DIR = Path("/app/sandbox/session_20251224_091022_74d424b00564/data")
RESULTS_DIR = Path("/app/sandbox/session_20251224_091022_74d424b00564/results")
RESULTS_DIR.mkdir(parents=True, exist_ok=True)

# Load unified dataset
print("Loading unified dataset...")
print("-" * 70)
df = pd.read_csv(DATA_DIR / "unified_dataset.csv")
print(f"✓ Loaded dataset: {df.shape}")
print(f"  States: {df['State'].nunique()}")
print(f"  Years: {df['Year'].min()}-{df['Year'].max()}")
print()

# ============================================================================
# PART 1: EXPLORATORY DATA ANALYSIS
# ============================================================================
print("=" * 70)
print("PART 1: EXPLORATORY DATA ANALYSIS")
print("=" * 70)
print()

# 1.1 Overall emigration statistics
print("1.1 Overall Emigration Statistics")
print("-" * 70)

total_emigration_by_year = df.groupby('Year')['Emigration_Count'].sum()
print("Total emigration by year:")
print(total_emigration_by_year)
print()

# Calculate growth rate
overall_growth = ((total_emigration_by_year.iloc[-1] - total_emigration_by_year.iloc[0]) /
                  total_emigration_by_year.iloc[0] * 100)
print(f"Overall growth 2010-2023: {overall_growth:.2f}%")
print()

# 1.2 Top states analysis
print("1.2 Top Emigrant States")
print("-" * 70)

# Top 10 states in 2023
top_states_2023 = df[df['Year'] == 2023].nlargest(10, 'Emigration_Count')[
    ['State', 'Emigration_Count', 'Emigration_Per_100K', 'Major_Destination']
]
print("Top 10 states (2023):")
print(top_states_2023)
print()

# Calculate concentration
total_2023 = df[df['Year'] == 2023]['Emigration_Count'].sum()
top5_2023 = df[df['Year'] == 2023].nlargest(5, 'Emigration_Count')['Emigration_Count'].sum()
concentration_ratio = (top5_2023 / total_2023) * 100

print(f"Top 5 states account for {concentration_ratio:.2f}% of total emigration (2023)")
print()

# 1.3 State comparison over time
print("1.3 Emigration Growth by State (2010-2023)")
print("-" * 70)

growth_by_state = df.groupby('State').apply(
    lambda x: ((x[x['Year'] == 2023]['Emigration_Count'].iloc[0] -
                x[x['Year'] == 2010]['Emigration_Count'].iloc[0]) /
               x[x['Year'] == 2010]['Emigration_Count'].iloc[0] * 100)
).sort_values(ascending=False)

print("Top 10 fastest growing states:")
print(growth_by_state.head(10))
print()
print("Slowest growing states:")
print(growth_by_state.tail(5))
print()

# ============================================================================
# PART 2: TREND ANALYSIS
# ============================================================================
print("=" * 70)
print("PART 2: TREND ANALYSIS")
print("=" * 70)
print()

# 2.1 State ranking changes
print("2.1 Changes in State Rankings")
print("-" * 70)

# Compare rankings in 2010 vs 2023
ranking_2010 = df[df['Year'] == 2010][['State', 'Emigration_Rank']].sort_values('Emigration_Rank')
ranking_2023 = df[df['Year'] == 2023][['State', 'Emigration_Rank']].sort_values('Emigration_Rank')

ranking_change = pd.merge(
    ranking_2010.rename(columns={'Emigration_Rank': 'Rank_2010'}),
    ranking_2023.rename(columns={'Emigration_Rank': 'Rank_2023'}),
    on='State'
)
ranking_change['Rank_Change'] = ranking_change['Rank_2010'] - ranking_change['Rank_2023']
ranking_change = ranking_change.sort_values('Rank_Change', ascending=False)

print("States with biggest ranking improvements (moved up):")
print(ranking_change.head(5))
print()
print("States with biggest ranking declines (moved down):")
print(ranking_change.tail(5))
print()

# 2.2 Temporal patterns by period
print("2.2 Average Emigration by Period")
print("-" * 70)

period_analysis = df.groupby(['Period', 'State'])['Emigration_Count'].mean().reset_index()
period_pivot = period_analysis.pivot(index='State', columns='Period', values='Emigration_Count')
period_pivot['Total_Growth'] = ((period_pivot['2020-2023'] - period_pivot['2010-2014']) /
                                period_pivot['2010-2014'] * 100)
period_pivot = period_pivot.sort_values('Total_Growth', ascending=False)

print("Growth across periods (top 10 states):")
print(period_pivot.head(10).round(2))
print()

# ============================================================================
# PART 3: CORRELATION ANALYSIS
# ============================================================================
print("=" * 70)
print("PART 3: CORRELATION ANALYSIS")
print("=" * 70)
print()

# Select variables for correlation analysis
correlation_vars = [
    'Emigration_Per_100K',
    'GDP_Per_Capita',
    'Unemployment_Rate',
    'Agriculture_Share_GSDP',
    'Literacy_Rate',
    'Higher_Education_Enrollment',
    'English_Proficiency_Index',
    'Remittances_Per_Capita',
    'Diaspora_Network_Strength'
]

# 3.1 Pearson correlations
print("3.1 Pearson Correlations with Emigration Rate")
print("-" * 70)

correlations = {}
p_values = {}

for var in correlation_vars[1:]:  # Skip emigration itself
    if var in df.columns:
        # Remove NaN values
        data_subset = df[['Emigration_Per_100K', var]].dropna()

        if len(data_subset) > 0:
            corr, p_val = pearsonr(data_subset['Emigration_Per_100K'], data_subset[var])
            correlations[var] = corr
            p_values[var] = p_val

# Sort by absolute correlation
corr_df = pd.DataFrame({
    'Variable': correlations.keys(),
    'Correlation': correlations.values(),
    'P_Value': p_values.values()
})
corr_df['Abs_Correlation'] = corr_df['Correlation'].abs()
corr_df = corr_df.sort_values('Abs_Correlation', ascending=False)

# Apply FDR correction (Benjamini-Hochberg)
from statsmodels.stats.multitest import multipletests
_, p_values_corrected, _, _ = multipletests(corr_df['P_Value'], method='fdr_bh')
corr_df['P_Value_FDR'] = p_values_corrected
corr_df['Significant_FDR'] = corr_df['P_Value_FDR'] < 0.05

print("Correlation with Emigration Rate (per 100K population):")
print(corr_df.to_string(index=False))
print()

# Interpret results
print("Interpretation:")
print("-" * 70)
for idx, row in corr_df.head(5).iterrows():
    direction = "positive" if row['Correlation'] > 0 else "negative"
    strength = "strong" if abs(row['Correlation']) > 0.5 else "moderate" if abs(row['Correlation']) > 0.3 else "weak"
    sig = "significant" if row['Significant_FDR'] else "not significant"
    print(f"  {row['Variable']}: {strength} {direction} correlation (r={row['Correlation']:.3f}, {sig})")
print()

# 3.2 Correlation matrix for all factors
print("3.2 Full Correlation Matrix")
print("-" * 70)

corr_matrix = df[correlation_vars].corr()
print(corr_matrix.round(3))
print()

# ============================================================================
# PART 4: REGRESSION ANALYSIS
# ============================================================================
print("=" * 70)
print("PART 4: REGRESSION ANALYSIS")
print("=" * 70)
print()

# Prepare data for regression (remove NaN values)
regression_data = df[correlation_vars].dropna()

print(f"Regression dataset: {regression_data.shape[0]} observations")
print()

# 4.1 Simple linear regressions
print("4.1 Simple Linear Regressions")
print("-" * 70)

simple_regression_results = []

for var in correlation_vars[1:]:
    X = sm.add_constant(regression_data[var])
    y = regression_data['Emigration_Per_100K']

    model = sm.OLS(y, X).fit()

    simple_regression_results.append({
        'Variable': var,
        'Coefficient': model.params[var],
        'Std_Error': model.bse[var],
        'T_Statistic': model.tvalues[var],
        'P_Value': model.pvalues[var],
        'R_Squared': model.rsquared,
        'Adj_R_Squared': model.rsquared_adj
    })

simple_reg_df = pd.DataFrame(simple_regression_results)
simple_reg_df['Abs_Coefficient'] = simple_reg_df['Coefficient'].abs()
simple_reg_df = simple_reg_df.sort_values('R_Squared', ascending=False)

print("Simple regression results (sorted by R²):")
print(simple_reg_df.to_string(index=False))
print()

# 4.2 Multiple regression with all factors
print("4.2 Multiple Regression Model")
print("-" * 70)

# Model 1: All economic factors
X_economic = sm.add_constant(regression_data[['GDP_Per_Capita', 'Unemployment_Rate', 'Agriculture_Share_GSDP']])
y = regression_data['Emigration_Per_100K']
model_economic = sm.OLS(y, X_economic).fit()

print("Model 1: Economic Factors Only")
print(model_economic.summary())
print()

# Model 2: All social factors
X_social = sm.add_constant(regression_data[['Literacy_Rate', 'Higher_Education_Enrollment', 'English_Proficiency_Index']])
model_social = sm.OLS(y, X_social).fit()

print("Model 2: Social Factors Only")
print(model_social.summary())
print()

# Model 3: Historical factors
X_historical = sm.add_constant(regression_data[['Diaspora_Network_Strength', 'Remittances_Per_Capita']])
model_historical = sm.OLS(y, X_historical).fit()

print("Model 3: Historical/Network Factors Only")
print(model_historical.summary())
print()

# Model 4: Full model with all factors
X_full = sm.add_constant(regression_data[correlation_vars[1:]])
model_full = sm.OLS(y, X_full).fit()

print("Model 4: Full Model (All Factors)")
print(model_full.summary())
print()

# Compare models
print("4.3 Model Comparison")
print("-" * 70)

model_comparison = pd.DataFrame({
    'Model': ['Economic', 'Social', 'Historical', 'Full'],
    'R_Squared': [model_economic.rsquared, model_social.rsquared,
                  model_historical.rsquared, model_full.rsquared],
    'Adj_R_Squared': [model_economic.rsquared_adj, model_social.rsquared_adj,
                      model_historical.rsquared_adj, model_full.rsquared_adj],
    'AIC': [model_economic.aic, model_social.aic,
            model_historical.aic, model_full.aic],
    'BIC': [model_economic.bic, model_social.bic,
            model_historical.bic, model_full.bic]
})

print(model_comparison.to_string(index=False))
print()

# ============================================================================
# PART 5: KEY FINDINGS SUMMARY
# ============================================================================
print("=" * 70)
print("PART 5: KEY FINDINGS SUMMARY")
print("=" * 70)
print()

# Identify strongest predictors
strongest_predictors = corr_df[corr_df['Significant_FDR'] == True].head(3)

print("STRONGEST PREDICTORS OF EMIGRATION:")
print("-" * 70)
for idx, row in strongest_predictors.iterrows():
    print(f"  {idx+1}. {row['Variable']}")
    print(f"     Correlation: {row['Correlation']:.3f} (p < {row['P_Value_FDR']:.4f})")
    print()

print("MODEL PERFORMANCE:")
print("-" * 70)
best_model = model_comparison.loc[model_comparison['Adj_R_Squared'].idxmax(), 'Model']
best_r2 = model_comparison['Adj_R_Squared'].max()
print(f"  Best model: {best_model} (Adjusted R² = {best_r2:.3f})")
print(f"  This model explains {best_r2*100:.1f}% of variation in emigration rates")
print()

# ============================================================================
# SAVE RESULTS
# ============================================================================
print("=" * 70)
print("SAVING RESULTS")
print("=" * 70)
print()

# Save correlation results
corr_file = RESULTS_DIR / "correlation_results.csv"
corr_df.to_csv(corr_file, index=False)
print(f"✓ Saved: {corr_file}")

# Save regression results
simple_reg_file = RESULTS_DIR / "simple_regression_results.csv"
simple_reg_df.to_csv(simple_reg_file, index=False)
print(f"✓ Saved: {simple_reg_file}")

# Save model comparison
model_comp_file = RESULTS_DIR / "model_comparison.csv"
model_comparison.to_csv(model_comp_file, index=False)
print(f"✓ Saved: {model_comp_file}")

# Save full correlation matrix
corr_matrix_file = RESULTS_DIR / "correlation_matrix.csv"
corr_matrix.to_csv(corr_matrix_file)
print(f"✓ Saved: {corr_matrix_file}")

# Save summary statistics
summary_stats = {
    'analysis_date': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
    'dataset_shape': df.shape,
    'years_covered': f"{df['Year'].min()}-{df['Year'].max()}",
    'states_analyzed': int(df['State'].nunique()),
    'overall_growth_2010_2023_pct': float(overall_growth),
    'top_5_concentration_2023_pct': float(concentration_ratio),
    'strongest_predictors': [
        {
            'variable': row['Variable'],
            'correlation': float(row['Correlation']),
            'p_value_fdr': float(row['P_Value_FDR'])
        }
        for _, row in strongest_predictors.iterrows()
    ],
    'best_model': {
        'name': best_model,
        'adjusted_r_squared': float(best_r2),
        'variance_explained_pct': float(best_r2 * 100)
    },
    'top_emigrant_states_2023': [
        {'state': row['State'], 'emigration_count': int(row['Emigration_Count'])}
        for _, row in top_states_2023.head(5).iterrows()
    ]
}

summary_file = RESULTS_DIR / "statistical_summary.json"
with open(summary_file, 'w') as f:
    json.dump(summary_stats, f, indent=2)
print(f"✓ Saved: {summary_file}")

# Save ranking changes
ranking_file = RESULTS_DIR / "state_ranking_changes.csv"
ranking_change.to_csv(ranking_file, index=False)
print(f"✓ Saved: {ranking_file}")

print()
print(f"End time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
print("=" * 70)
print("✓ STATISTICAL ANALYSIS COMPLETE")
print("=" * 70)