"""
Visualization Script for Indian Emigration Analysis

Creates publication-quality visualizations for slides:
1. Time series trends
2. State comparisons
3. Correlation heatmaps
4. Regression coefficient plots
5. Factor relationship scatter plots
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path
from datetime import datetime

# Set style for publication-quality figures
plt.style.use('seaborn-v0_8-darkgrid')
sns.set_palette("husl")
plt.rcParams['figure.dpi'] = 300
plt.rcParams['savefig.dpi'] = 300
plt.rcParams['font.family'] = 'sans-serif'
plt.rcParams['font.size'] = 10
plt.rcParams['axes.labelsize'] = 11
plt.rcParams['axes.titlesize'] = 12
plt.rcParams['xtick.labelsize'] = 9
plt.rcParams['ytick.labelsize'] = 9
plt.rcParams['legend.fontsize'] = 9

print("=" * 70)
print("CREATING VISUALIZATIONS FOR SLIDES")
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")
FIGURES_DIR = Path("/app/sandbox/session_20251224_091022_74d424b00564/figures")
FIGURES_DIR.mkdir(parents=True, exist_ok=True)

# Load data
print("Loading datasets...")
df = pd.read_csv(DATA_DIR / "unified_dataset.csv")
corr_results = pd.read_csv(RESULTS_DIR / "correlation_results.csv")
corr_matrix = pd.read_csv(RESULTS_DIR / "correlation_matrix.csv", index_col=0)

print(f"✓ Loaded data: {df.shape}")
print()

# ============================================================================
# FIGURE 1: Overall Emigration Trend (2010-2023)
# ============================================================================
print("Creating Figure 1: Overall emigration trend...")

fig, ax = plt.subplots(figsize=(10, 6))

total_by_year = df.groupby('Year')['Emigration_Count'].sum() / 1_000_000

ax.plot(total_by_year.index, total_by_year.values, marker='o', linewidth=2.5,
        markersize=8, color='#2E86AB')
ax.fill_between(total_by_year.index, 0, total_by_year.values, alpha=0.3, color='#2E86AB')

ax.set_xlabel('Year', fontsize=12, fontweight='bold')
ax.set_ylabel('Total Emigration (Millions)', fontsize=12, fontweight='bold')
ax.set_title('Total Emigration from India (2010-2023)', fontsize=14, fontweight='bold', pad=20)
ax.grid(True, alpha=0.3, linestyle='--')

# Add growth annotation
growth_text = f"Growth: +151.6%\n(2010-2023)"
ax.text(0.05, 0.95, growth_text, transform=ax.transAxes,
        fontsize=11, fontweight='bold',
        verticalalignment='top',
        bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))

plt.tight_layout()
plt.savefig(FIGURES_DIR / "01_overall_emigration_trend.png", dpi=300, bbox_inches='tight')
plt.close()
print("✓ Saved: 01_overall_emigration_trend.png")

# ============================================================================
# FIGURE 2: Top 10 Emigrant States (2023)
# ============================================================================
print("Creating Figure 2: Top 10 emigrant states...")

top_10_2023 = df[df['Year'] == 2023].nlargest(10, 'Emigration_Count')

fig, ax = plt.subplots(figsize=(10, 6))

bars = ax.barh(range(len(top_10_2023)), top_10_2023['Emigration_Count'] / 1000,
               color=sns.color_palette("viridis", len(top_10_2023)))

ax.set_yticks(range(len(top_10_2023)))
ax.set_yticklabels(top_10_2023['State'])
ax.set_xlabel('Emigration Count (Thousands)', fontsize=12, fontweight='bold')
ax.set_title('Top 10 Emigrant States in 2023', fontsize=14, fontweight='bold', pad=20)
ax.invert_yaxis()
ax.grid(True, axis='x', alpha=0.3, linestyle='--')

# Add value labels
for i, (idx, row) in enumerate(top_10_2023.iterrows()):
    value = row['Emigration_Count'] / 1000
    ax.text(value + 5, i, f"{value:.0f}K", va='center', fontweight='bold')

plt.tight_layout()
plt.savefig(FIGURES_DIR / "02_top_10_states_2023.png", dpi=300, bbox_inches='tight')
plt.close()
print("✓ Saved: 02_top_10_states_2023.png")

# ============================================================================
# FIGURE 3: Emigration Trends by Top 7 States
# ============================================================================
print("Creating Figure 3: Emigration trends by top states...")

# Get top 7 states by average emigration
top_states = df.groupby('State')['Emigration_Count'].mean().nlargest(7).index

fig, ax = plt.subplots(figsize=(12, 7))

for state in top_states:
    state_data = df[df['State'] == state].sort_values('Year')
    ax.plot(state_data['Year'], state_data['Emigration_Count'] / 1000,
            marker='o', linewidth=2.5, markersize=6, label=state)

ax.set_xlabel('Year', fontsize=12, fontweight='bold')
ax.set_ylabel('Emigration Count (Thousands)', fontsize=12, fontweight='bold')
ax.set_title('Emigration Trends: Top 7 States (2010-2023)', fontsize=14, fontweight='bold', pad=20)
ax.legend(loc='upper left', framealpha=0.9)
ax.grid(True, alpha=0.3, linestyle='--')

plt.tight_layout()
plt.savefig(FIGURES_DIR / "03_emigration_trends_top_states.png", dpi=300, bbox_inches='tight')
plt.close()
print("✓ Saved: 03_emigration_trends_top_states.png")

# ============================================================================
# FIGURE 4: State Ranking Changes (2010 vs 2023)
# ============================================================================
print("Creating Figure 4: State ranking changes...")

ranking_changes = pd.read_csv(RESULTS_DIR / "state_ranking_changes.csv")
ranking_changes = ranking_changes.sort_values('Rank_Change', ascending=False)

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

colors = ['green' if x > 0 else 'red' if x < 0 else 'gray'
          for x in ranking_changes['Rank_Change']]

bars = ax.barh(range(len(ranking_changes)), ranking_changes['Rank_Change'], color=colors)

ax.set_yticks(range(len(ranking_changes)))
ax.set_yticklabels(ranking_changes['State'])
ax.set_xlabel('Rank Change (Positive = Moved Up)', fontsize=12, fontweight='bold')
ax.set_title('Change in State Rankings: 2010 to 2023', fontsize=14, fontweight='bold', pad=20)
ax.axvline(x=0, color='black', linestyle='-', linewidth=1)
ax.grid(True, axis='x', alpha=0.3, linestyle='--')

plt.tight_layout()
plt.savefig(FIGURES_DIR / "04_state_ranking_changes.png", dpi=300, bbox_inches='tight')
plt.close()
print("✓ Saved: 04_state_ranking_changes.png")

# ============================================================================
# FIGURE 5: Correlation Heatmap
# ============================================================================
print("Creating Figure 5: Correlation heatmap...")

# Select key variables for heatmap
vars_for_heatmap = [
    'Emigration_Per_100K',
    'GDP_Per_Capita',
    'Unemployment_Rate',
    'Agriculture_Share_GSDP',
    'Literacy_Rate',
    'English_Proficiency_Index',
    'Remittances_Per_Capita',
    'Diaspora_Network_Strength'
]

# Create nicer labels
label_map = {
    'Emigration_Per_100K': 'Emigration Rate',
    'GDP_Per_Capita': 'GDP per Capita',
    'Unemployment_Rate': 'Unemployment',
    'Agriculture_Share_GSDP': 'Agriculture Share',
    'Literacy_Rate': 'Literacy Rate',
    'English_Proficiency_Index': 'English Proficiency',
    'Remittances_Per_Capita': 'Remittances',
    'Diaspora_Network_Strength': 'Diaspora Network'
}

corr_subset = corr_matrix.loc[vars_for_heatmap, vars_for_heatmap]
corr_subset = corr_subset.rename(index=label_map, columns=label_map)

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

mask = np.triu(np.ones_like(corr_subset, dtype=bool), k=1)
sns.heatmap(corr_subset, 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: Emigration and Key Factors', fontsize=14, fontweight='bold', pad=20)

plt.tight_layout()
plt.savefig(FIGURES_DIR / "05_correlation_heatmap.png", dpi=300, bbox_inches='tight')
plt.close()
print("✓ Saved: 05_correlation_heatmap.png")

# ============================================================================
# FIGURE 6: Key Predictors - Correlation Coefficients
# ============================================================================
print("Creating Figure 6: Key predictors bar chart...")

# Get top 5 significant predictors
top_predictors = corr_results[corr_results['Significant_FDR'] == True].head(5)

fig, ax = plt.subplots(figsize=(10, 6))

colors = ['#2E86AB' if x > 0 else '#A23B72' for x in top_predictors['Correlation']]

bars = ax.barh(range(len(top_predictors)), top_predictors['Correlation'], color=colors)

ax.set_yticks(range(len(top_predictors)))
ax.set_yticklabels([label_map.get(v, v) for v in top_predictors['Variable']])
ax.set_xlabel('Pearson Correlation Coefficient', fontsize=12, fontweight='bold')
ax.set_title('Top 5 Predictors of Emigration Rate', fontsize=14, fontweight='bold', pad=20)
ax.axvline(x=0, color='black', linestyle='-', linewidth=1)
ax.set_xlim(-0.3, 0.8)
ax.grid(True, axis='x', alpha=0.3, linestyle='--')

# Add value labels
for i, (idx, row) in enumerate(top_predictors.iterrows()):
    value = row['Correlation']
    ax.text(value + 0.02 if value > 0 else value - 0.02, i,
            f"r = {value:.3f}", va='center', ha='left' if value > 0 else 'right',
            fontweight='bold')

# Add significance note
ax.text(0.98, 0.02, 'All correlations significant (FDR < 0.05)',
        transform=ax.transAxes, ha='right', va='bottom',
        fontsize=9, style='italic',
        bbox=dict(boxstyle='round', facecolor='yellow', alpha=0.3))

plt.tight_layout()
plt.savefig(FIGURES_DIR / "06_key_predictors_correlation.png", dpi=300, bbox_inches='tight')
plt.close()
print("✓ Saved: 06_key_predictors_correlation.png")

# ============================================================================
# FIGURE 7: Scatter - Diaspora Network vs Emigration
# ============================================================================
print("Creating Figure 7: Diaspora network vs emigration scatter...")

fig, ax = plt.subplots(figsize=(10, 7))

# Use 2023 data for clarity
df_2023 = df[df['Year'] == 2023]

scatter = ax.scatter(df_2023['Diaspora_Network_Strength'],
                    df_2023['Emigration_Per_100K'],
                    s=df_2023['Emigration_Count'] / 1000,
                    c=df_2023['Emigration_Per_100K'],
                    cmap='viridis', alpha=0.7, edgecolors='black', linewidth=1)

# Add trend line
z = np.polyfit(df_2023['Diaspora_Network_Strength'], df_2023['Emigration_Per_100K'], 1)
p = np.poly1d(z)
x_line = np.linspace(df_2023['Diaspora_Network_Strength'].min(),
                     df_2023['Diaspora_Network_Strength'].max(), 100)
ax.plot(x_line, p(x_line), "r--", linewidth=2, label='Trend line')

# Label top states
top_5_states = df_2023.nlargest(5, 'Emigration_Count')
for idx, row in top_5_states.iterrows():
    ax.annotate(row['State'],
                (row['Diaspora_Network_Strength'], row['Emigration_Per_100K']),
                xytext=(5, 5), textcoords='offset points',
                fontsize=8, fontweight='bold',
                bbox=dict(boxstyle='round,pad=0.3', facecolor='yellow', alpha=0.7))

ax.set_xlabel('Diaspora Network Strength Index', fontsize=12, fontweight='bold')
ax.set_ylabel('Emigration Rate (per 100K population)', fontsize=12, fontweight='bold')
ax.set_title('Diaspora Networks Drive Emigration (2023)', fontsize=14, fontweight='bold', pad=20)
ax.legend(loc='upper left')
ax.grid(True, alpha=0.3, linestyle='--')

# Add colorbar
cbar = plt.colorbar(scatter, ax=ax)
cbar.set_label('Emigration Rate', rotation=270, labelpad=20)

# Add correlation text
corr_value = corr_results[corr_results['Variable'] == 'Diaspora_Network_Strength']['Correlation'].values[0]
ax.text(0.02, 0.98, f'Correlation: r = {corr_value:.3f}',
        transform=ax.transAxes, fontsize=11, fontweight='bold',
        verticalalignment='top',
        bbox=dict(boxstyle='round', facecolor='white', alpha=0.9))

plt.tight_layout()
plt.savefig(FIGURES_DIR / "07_diaspora_network_scatter.png", dpi=300, bbox_inches='tight')
plt.close()
print("✓ Saved: 07_diaspora_network_scatter.png")

# ============================================================================
# FIGURE 8: Scatter - Remittances vs Emigration
# ============================================================================
print("Creating Figure 8: Remittances vs emigration scatter...")

fig, ax = plt.subplots(figsize=(10, 7))

scatter = ax.scatter(df_2023['Remittances_Per_Capita'],
                    df_2023['Emigration_Per_100K'],
                    s=df_2023['Emigration_Count'] / 1000,
                    c=df_2023['GDP_Per_Capita'],
                    cmap='coolwarm', alpha=0.7, edgecolors='black', linewidth=1)

# Add trend line
z = np.polyfit(df_2023['Remittances_Per_Capita'], df_2023['Emigration_Per_100K'], 1)
p = np.poly1d(z)
x_line = np.linspace(df_2023['Remittances_Per_Capita'].min(),
                     df_2023['Remittances_Per_Capita'].max(), 100)
ax.plot(x_line, p(x_line), "r--", linewidth=2, label='Trend line')

# Label top states
for idx, row in top_5_states.iterrows():
    ax.annotate(row['State'],
                (row['Remittances_Per_Capita'], row['Emigration_Per_100K']),
                xytext=(5, 5), textcoords='offset points',
                fontsize=8, fontweight='bold',
                bbox=dict(boxstyle='round,pad=0.3', facecolor='yellow', alpha=0.7))

ax.set_xlabel('Remittances per Capita (USD)', fontsize=12, fontweight='bold')
ax.set_ylabel('Emigration Rate (per 100K population)', fontsize=12, fontweight='bold')
ax.set_title('Remittances and Emigration (2023)', fontsize=14, fontweight='bold', pad=20)
ax.legend(loc='upper left')
ax.grid(True, alpha=0.3, linestyle='--')

# Add colorbar
cbar = plt.colorbar(scatter, ax=ax)
cbar.set_label('GDP per Capita (INR)', rotation=270, labelpad=20)

# Add correlation text
corr_value = corr_results[corr_results['Variable'] == 'Remittances_Per_Capita']['Correlation'].values[0]
ax.text(0.02, 0.98, f'Correlation: r = {corr_value:.3f}\n(Strongest predictor)',
        transform=ax.transAxes, fontsize=11, fontweight='bold',
        verticalalignment='top',
        bbox=dict(boxstyle='round', facecolor='white', alpha=0.9))

plt.tight_layout()
plt.savefig(FIGURES_DIR / "08_remittances_scatter.png", dpi=300, bbox_inches='tight')
plt.close()
print("✓ Saved: 08_remittances_scatter.png")

# ============================================================================
# FIGURE 9: Unemployment vs Emigration by State
# ============================================================================
print("Creating Figure 9: Unemployment vs emigration...")

fig, ax = plt.subplots(figsize=(10, 7))

scatter = ax.scatter(df_2023['Unemployment_Rate'],
                    df_2023['Emigration_Per_100K'],
                    s=df_2023['Emigration_Count'] / 1000,
                    c=df_2023['Literacy_Rate'],
                    cmap='RdYlGn', alpha=0.7, edgecolors='black', linewidth=1)

# Add trend line
z = np.polyfit(df_2023['Unemployment_Rate'], df_2023['Emigration_Per_100K'], 1)
p = np.poly1d(z)
x_line = np.linspace(df_2023['Unemployment_Rate'].min(),
                     df_2023['Unemployment_Rate'].max(), 100)
ax.plot(x_line, p(x_line), "r--", linewidth=2, label='Trend line')

# Label interesting states
states_to_label = ['Kerala', 'Punjab', 'Gujarat', 'Bihar', 'Uttar Pradesh']
for state in states_to_label:
    row = df_2023[df_2023['State'] == state].iloc[0]
    ax.annotate(state,
                (row['Unemployment_Rate'], row['Emigration_Per_100K']),
                xytext=(5, 5), textcoords='offset points',
                fontsize=8, fontweight='bold',
                bbox=dict(boxstyle='round,pad=0.3', facecolor='lightblue', alpha=0.7))

ax.set_xlabel('Unemployment Rate (%)', fontsize=12, fontweight='bold')
ax.set_ylabel('Emigration Rate (per 100K population)', fontsize=12, fontweight='bold')
ax.set_title('Unemployment and Emigration (2023)', fontsize=14, fontweight='bold', pad=20)
ax.legend(loc='upper right')
ax.grid(True, alpha=0.3, linestyle='--')

# Add colorbar
cbar = plt.colorbar(scatter, ax=ax)
cbar.set_label('Literacy Rate (%)', rotation=270, labelpad=20)

# Add correlation text
corr_value = corr_results[corr_results['Variable'] == 'Unemployment_Rate']['Correlation'].values[0]
ax.text(0.02, 0.98, f'Correlation: r = {corr_value:.3f}',
        transform=ax.transAxes, fontsize=11, fontweight='bold',
        verticalalignment='top',
        bbox=dict(boxstyle='round', facecolor='white', alpha=0.9))

plt.tight_layout()
plt.savefig(FIGURES_DIR / "09_unemployment_scatter.png", dpi=300, bbox_inches='tight')
plt.close()
print("✓ Saved: 09_unemployment_scatter.png")

# ============================================================================
# FIGURE 10: Model Comparison (R²)
# ============================================================================
print("Creating Figure 10: Model comparison...")

model_comparison = pd.read_csv(RESULTS_DIR / "model_comparison.csv")

fig, ax = plt.subplots(figsize=(10, 6))

x_pos = range(len(model_comparison))
bars = ax.bar(x_pos, model_comparison['Adj_R_Squared'],
              color=['#FF6B6B', '#4ECDC4', '#45B7D1', '#96CEB4'])

ax.set_xticks(x_pos)
ax.set_xticklabels(model_comparison['Model'], fontsize=11, fontweight='bold')
ax.set_ylabel('Adjusted R² (Variance Explained)', fontsize=12, fontweight='bold')
ax.set_title('Model Comparison: Which Factors Explain Emigration?', fontsize=14, fontweight='bold', pad=20)
ax.set_ylim(0, 1.0)
ax.grid(True, axis='y', alpha=0.3, linestyle='--')

# Add value labels and percentage
for i, (idx, row) in enumerate(model_comparison.iterrows()):
    value = row['Adj_R_Squared']
    ax.text(i, value + 0.02, f"{value:.3f}\n({value*100:.1f}%)",
            ha='center', fontweight='bold', fontsize=10)

# Highlight the best model
best_idx = model_comparison['Adj_R_Squared'].idxmax()
ax.patches[best_idx].set_edgecolor('gold')
ax.patches[best_idx].set_linewidth(3)

plt.tight_layout()
plt.savefig(FIGURES_DIR / "10_model_comparison.png", dpi=300, bbox_inches='tight')
plt.close()
print("✓ Saved: 10_model_comparison.png")

# ============================================================================
# FIGURE 11: Top vs Bottom States Comparison
# ============================================================================
print("Creating Figure 11: Top vs bottom states comparison...")

# Get top 5 and bottom 5 states by average emigration
state_avg = df.groupby('State')['Emigration_Per_100K'].mean().sort_values(ascending=False)
top_5_avg = state_avg.head(5)
bottom_5_avg = state_avg.tail(5)

# Get their characteristics (2023 data)
comparison_states = list(top_5_avg.index) + list(bottom_5_avg.index)
comparison_data = df[df['Year'] == 2023][df['State'].isin(comparison_states)]

# Create comparison dataframe
comparison_metrics = []
for state in comparison_states:
    state_data = comparison_data[comparison_data['State'] == state].iloc[0]
    group = 'High Emigration' if state in top_5_avg.index else 'Low Emigration'
    comparison_metrics.append({
        'State': state,
        'Group': group,
        'GDP_Per_Capita': state_data['GDP_Per_Capita'] / 1000,  # in thousands
        'Unemployment': state_data['Unemployment_Rate'],
        'Literacy': state_data['Literacy_Rate'],
        'Diaspora': state_data['Diaspora_Network_Strength']
    })

comp_df = pd.DataFrame(comparison_metrics)

# Create subplots
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
fig.suptitle('High vs Low Emigration States: Key Characteristics (2023)',
             fontsize=16, fontweight='bold', y=1.00)

metrics = [
    ('GDP_Per_Capita', 'GDP per Capita (₹ thousands)', 0),
    ('Unemployment', 'Unemployment Rate (%)', 1),
    ('Literacy', 'Literacy Rate (%)', 2),
    ('Diaspora', 'Diaspora Network Strength', 3)
]

for metric, label, idx in metrics:
    ax = axes[idx // 2, idx % 2]

    # Separate data by group
    high_emig = comp_df[comp_df['Group'] == 'High Emigration'][metric]
    low_emig = comp_df[comp_df['Group'] == 'Low Emigration'][metric]

    # Create box plot
    bp = ax.boxplot([high_emig, low_emig], labels=['High\nEmigration', 'Low\nEmigration'],
                     patch_artist=True, showmeans=True)

    # Color boxes
    bp['boxes'][0].set_facecolor('#FF6B6B')
    bp['boxes'][1].set_facecolor('#4ECDC4')

    ax.set_ylabel(label, fontsize=11, fontweight='bold')
    ax.set_title(label, fontsize=12, fontweight='bold')
    ax.grid(True, axis='y', alpha=0.3, linestyle='--')

    # Add mean values
    ax.text(1, high_emig.mean(), f"μ={high_emig.mean():.1f}",
            ha='center', va='bottom', fontweight='bold', fontsize=9)
    ax.text(2, low_emig.mean(), f"μ={low_emig.mean():.1f}",
            ha='center', va='bottom', fontweight='bold', fontsize=9)

plt.tight_layout()
plt.savefig(FIGURES_DIR / "11_high_vs_low_emigration_states.png", dpi=300, bbox_inches='tight')
plt.close()
print("✓ Saved: 11_high_vs_low_emigration_states.png")

# ============================================================================
# Summary
# ============================================================================
print()
print("=" * 70)
print("VISUALIZATION SUMMARY")
print("=" * 70)
print(f"Created 11 publication-quality figures:")
print("  1. Overall emigration trend (2010-2023)")
print("  2. Top 10 emigrant states (2023)")
print("  3. Emigration trends by top 7 states")
print("  4. State ranking changes (2010 vs 2023)")
print("  5. Correlation heatmap")
print("  6. Key predictors (correlation coefficients)")
print("  7. Diaspora network vs emigration scatter")
print("  8. Remittances vs emigration scatter")
print("  9. Unemployment vs emigration scatter")
print(" 10. Model comparison (R² values)")
print(" 11. High vs low emigration states comparison")
print()
print(f"All figures saved to: {FIGURES_DIR}")
print(f"Resolution: 300 DPI (publication quality)")
print()
print(f"End time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
print("=" * 70)
print("✓ VISUALIZATION CREATION COMPLETE")
print("=" * 70)