"""
Data Processing and Integration Script

Merges all acquired datasets into a unified analytical format.
Performs data cleaning, validation, and feature engineering.
"""

import pandas as pd
import numpy as np
from pathlib import Path
from datetime import datetime

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

# Define paths
RAW_DIR = Path("/app/sandbox/session_20251224_091022_74d424b00564/data/raw")
PROCESSED_DIR = Path("/app/sandbox/session_20251224_091022_74d424b00564/data")
PROCESSED_DIR.mkdir(parents=True, exist_ok=True)

# Step 1: Load all datasets
print("Step 1: Loading raw datasets...")
print("-" * 70)

emigration_df = pd.read_csv(RAW_DIR / "state_emigration_data.csv")
economic_df = pd.read_csv(RAW_DIR / "state_economic_indicators.csv")
social_df = pd.read_csv(RAW_DIR / "state_social_indicators.csv")
historical_df = pd.read_csv(RAW_DIR / "state_historical_diaspora.csv")

print(f"✓ Loaded emigration data: {emigration_df.shape}")
print(f"✓ Loaded economic data: {economic_df.shape}")
print(f"✓ Loaded social data: {social_df.shape}")
print(f"✓ Loaded historical data: {historical_df.shape}")
print()

# Step 2: Data validation
print("Step 2: Validating data integrity...")
print("-" * 70)

# Check for missing values
print("Missing values check:")
print(f"  Emigration: {emigration_df.isnull().sum().sum()} missing")
print(f"  Economic: {economic_df.isnull().sum().sum()} missing")
print(f"  Social: {social_df.isnull().sum().sum()} missing")
print(f"  Historical: {historical_df.isnull().sum().sum()} missing")
print()

# Check state-year coverage
states = emigration_df['State'].unique()
years = emigration_df['Year'].unique()
print(f"Coverage: {len(states)} states × {len(years)} years = {len(states) * len(years)} expected records")
print(f"  Emigration: {len(emigration_df)} records")
print(f"  Economic: {len(economic_df)} records")
print(f"  Social: {len(social_df)} records")
print(f"  Historical: {len(historical_df)} records")
print()

# Verify all datasets have same state-year combinations
assert len(emigration_df) == len(economic_df) == len(social_df) == len(historical_df), \
    "Datasets have different lengths!"
print("✓ All datasets have consistent coverage")
print()

# Step 3: Merge datasets
print("Step 3: Merging datasets...")
print("-" * 70)

# Start with emigration data
unified_df = emigration_df.copy()

# Merge economic indicators
unified_df = unified_df.merge(
    economic_df,
    on=['State', 'Year'],
    how='left',
    validate='one_to_one'
)

# Merge social indicators
unified_df = unified_df.merge(
    social_df,
    on=['State', 'Year'],
    how='left',
    validate='one_to_one'
)

# Merge historical/diaspora data
unified_df = unified_df.merge(
    historical_df,
    on=['State', 'Year'],
    how='left',
    validate='one_to_one'
)

print(f"✓ Merged dataset shape: {unified_df.shape}")
print(f"✓ Columns: {list(unified_df.columns)}")
print()

# Check for merge issues
if unified_df.isnull().sum().sum() > 0:
    print("⚠ Warning: Merge produced missing values:")
    print(unified_df.isnull().sum()[unified_df.isnull().sum() > 0])
else:
    print("✓ No missing values after merge")
print()

# Step 4: Feature engineering
print("Step 4: Creating derived features...")
print("-" * 70)

# Calculate population estimates (using back-calculation from GDP per capita)
# India's average GDP per capita ~$2,700 (2023), population ~1.4 billion
# This gives us GSDP = GDP_per_capita * population
unified_df['Population_Estimate'] = (
    unified_df['GDP_Per_Capita'] * 1000  # Rough population estimate
).round(0).astype(int)

# Emigration rate per 100,000 population
unified_df['Emigration_Per_100K'] = (
    (unified_df['Emigration_Count'] / unified_df['Population_Estimate']) * 100000
).round(2)

# Remittances per capita
unified_df['Remittances_Per_Capita'] = (
    (unified_df['Remittances_USD_Millions'] * 1_000_000) / unified_df['Population_Estimate']
).round(2)

# Create categorical variables
# Emigration intensity category
unified_df['Emigration_Category'] = pd.cut(
    unified_df['Emigration_Per_100K'],
    bins=[0, 50, 100, 200, float('inf')],
    labels=['Low', 'Medium', 'High', 'Very High']
)

# Economic development category based on GDP per capita
unified_df['Development_Category'] = pd.cut(
    unified_df['GDP_Per_Capita'],
    bins=[0, 150000, 250000, 350000, float('inf')],
    labels=['Low Income', 'Lower-Middle', 'Upper-Middle', 'High Income']
)

# Create time period categories
unified_df['Period'] = pd.cut(
    unified_df['Year'],
    bins=[2009, 2014, 2019, 2024],
    labels=['2010-2014', '2015-2019', '2020-2023']
)

print("✓ Created derived features:")
print("  - Population_Estimate")
print("  - Emigration_Per_100K")
print("  - Remittances_Per_Capita")
print("  - Emigration_Category")
print("  - Development_Category")
print("  - Period")
print()

# Step 5: Calculate growth rates and changes
print("Step 5: Calculating temporal changes...")
print("-" * 70)

# Sort by state and year
unified_df = unified_df.sort_values(['State', 'Year']).reset_index(drop=True)

# Calculate year-over-year growth rates for key variables
for var in ['Emigration_Count', 'GDP_Per_Capita', 'Unemployment_Rate', 'Literacy_Rate']:
    unified_df[f'{var}_Growth'] = unified_df.groupby('State')[var].pct_change() * 100

# Calculate cumulative change since 2010
for state in states:
    state_mask = unified_df['State'] == state
    base_emigration = unified_df[state_mask & (unified_df['Year'] == 2010)]['Emigration_Count'].values[0]

    unified_df.loc[state_mask, 'Emigration_Change_Since_2010'] = (
        (unified_df.loc[state_mask, 'Emigration_Count'] - base_emigration) / base_emigration * 100
    ).round(2)

print("✓ Calculated growth rates and temporal changes")
print()

# Step 6: Add state rankings
print("Step 6: Adding state rankings...")
print("-" * 70)

# Rank states by emigration count within each year
unified_df['Emigration_Rank'] = unified_df.groupby('Year')['Emigration_Count'].rank(
    method='dense',
    ascending=False
).astype(int)

# Rank by GDP per capita
unified_df['GDP_Rank'] = unified_df.groupby('Year')['GDP_Per_Capita'].rank(
    method='dense',
    ascending=False
).astype(int)

print("✓ Added state rankings")
print()

# Step 7: Data quality summary
print("Step 7: Final data quality check...")
print("-" * 70)

print(f"Final dataset shape: {unified_df.shape}")
print(f"Date range: {unified_df['Year'].min()} - {unified_df['Year'].max()}")
print(f"States covered: {unified_df['State'].nunique()}")
print()

print("Summary statistics:")
print(unified_df[['Emigration_Count', 'GDP_Per_Capita', 'Unemployment_Rate',
                  'Literacy_Rate', 'Emigration_Per_100K']].describe().round(2))
print()

# Check data types
print("Data types:")
print(unified_df.dtypes)
print()

# Save processed dataset
print("Step 8: Saving processed dataset...")
print("-" * 70)

output_file = PROCESSED_DIR / "unified_dataset.csv"
unified_df.to_csv(output_file, index=False)
print(f"✓ Saved: {output_file}")
print()

# Create a summary by state
state_summary = unified_df.groupby('State').agg({
    'Emigration_Count': ['mean', 'min', 'max'],
    'GDP_Per_Capita': 'mean',
    'Unemployment_Rate': 'mean',
    'Literacy_Rate': 'mean',
    'Diaspora_Network_Strength': 'mean',
    'Major_Destination': 'first'
}).round(2)

state_summary.columns = ['_'.join(col).strip() for col in state_summary.columns.values]
state_summary = state_summary.reset_index()
state_summary = state_summary.sort_values('Emigration_Count_mean', ascending=False)

summary_file = PROCESSED_DIR / "state_summary.csv"
state_summary.to_csv(summary_file, index=False)
print(f"✓ Saved state summary: {summary_file}")
print()

# Display top emigrant states
print("Top 10 emigrant states (average 2010-2023):")
print(state_summary[['State', 'Emigration_Count_mean', 'Major_Destination_first']].head(10))
print()

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