src/data_generation.py

"""
TB Vulnerability Hotspot Prediction - Synthetic Data Generation
================================================================
Generates realistic district-level Indian TB vulnerability data based on 
published distributions from NFHS-5, Census 2011, WHO GTB, and NIKSHAY.

Feature categories:
1. Demographic (Census 2011 distributions)
2. Nutritional/Health (NFHS-5 distributions)
3. Socioeconomic (SECC/Census)
4. Environmental (ERA5/CPCB distributions)
5. Healthcare Infrastructure (HMIS/NFHS)
6. Spatial (lat/lon of district centroids)
"""

import numpy as np
import pandas as pd
from scipy import stats

np.random.seed(42)

# ============================================================
# INDIAN DISTRICTS - Realistic centroids for ~640 districts
# ============================================================
# Based on actual Indian state-wise district counts and approximate centroid ranges

STATES_CONFIG = {
    'Andhra Pradesh': {'n_districts': 13, 'lat': (13.5, 18.5), 'lon': (77.0, 84.0), 'region': 'South'},
    'Arunachal Pradesh': {'n_districts': 25, 'lat': (26.5, 29.0), 'lon': (91.5, 97.5), 'region': 'NorthEast'},
    'Assam': {'n_districts': 33, 'lat': (24.0, 27.5), 'lon': (89.5, 96.0), 'region': 'NorthEast'},
    'Bihar': {'n_districts': 38, 'lat': (24.0, 27.5), 'lon': (83.5, 88.0), 'region': 'East'},
    'Chhattisgarh': {'n_districts': 27, 'lat': (18.5, 24.0), 'lon': (80.5, 84.0), 'region': 'Central'},
    'Goa': {'n_districts': 2, 'lat': (15.0, 15.8), 'lon': (73.5, 74.3), 'region': 'West'},
    'Gujarat': {'n_districts': 33, 'lat': (20.5, 24.5), 'lon': (68.5, 74.5), 'region': 'West'},
    'Haryana': {'n_districts': 22, 'lat': (27.5, 30.5), 'lon': (74.5, 77.5), 'region': 'North'},
    'Himachal Pradesh': {'n_districts': 12, 'lat': (30.5, 33.5), 'lon': (75.5, 79.0), 'region': 'North'},
    'Jharkhand': {'n_districts': 24, 'lat': (22.0, 25.5), 'lon': (83.5, 87.5), 'region': 'East'},
    'Karnataka': {'n_districts': 30, 'lat': (11.5, 18.5), 'lon': (74.0, 78.5), 'region': 'South'},
    'Kerala': {'n_districts': 14, 'lat': (8.0, 12.5), 'lon': (75.0, 77.5), 'region': 'South'},
    'Madhya Pradesh': {'n_districts': 52, 'lat': (21.0, 26.5), 'lon': (74.0, 82.5), 'region': 'Central'},
    'Maharashtra': {'n_districts': 36, 'lat': (15.5, 22.0), 'lon': (72.5, 80.5), 'region': 'West'},
    'Manipur': {'n_districts': 16, 'lat': (24.0, 26.0), 'lon': (93.0, 94.5), 'region': 'NorthEast'},
    'Meghalaya': {'n_districts': 11, 'lat': (25.0, 26.0), 'lon': (89.5, 92.5), 'region': 'NorthEast'},
    'Mizoram': {'n_districts': 8, 'lat': (21.5, 24.5), 'lon': (92.0, 93.5), 'region': 'NorthEast'},
    'Nagaland': {'n_districts': 11, 'lat': (25.0, 27.0), 'lon': (93.5, 95.5), 'region': 'NorthEast'},
    'Odisha': {'n_districts': 30, 'lat': (18.0, 22.5), 'lon': (81.5, 87.5), 'region': 'East'},
    'Punjab': {'n_districts': 22, 'lat': (29.5, 32.5), 'lon': (73.5, 76.5), 'region': 'North'},
    'Rajasthan': {'n_districts': 33, 'lat': (23.0, 30.0), 'lon': (69.0, 78.0), 'region': 'North'},
    'Sikkim': {'n_districts': 4, 'lat': (27.0, 28.0), 'lon': (88.0, 89.0), 'region': 'NorthEast'},
    'Tamil Nadu': {'n_districts': 38, 'lat': (8.0, 13.5), 'lon': (76.0, 80.5), 'region': 'South'},
    'Telangana': {'n_districts': 33, 'lat': (15.5, 19.5), 'lon': (77.0, 81.0), 'region': 'South'},
    'Tripura': {'n_districts': 8, 'lat': (22.5, 24.5), 'lon': (91.0, 92.5), 'region': 'NorthEast'},
    'Uttar Pradesh': {'n_districts': 75, 'lat': (23.5, 30.5), 'lon': (77.0, 84.5), 'region': 'North'},
    'Uttarakhand': {'n_districts': 13, 'lat': (28.5, 31.5), 'lon': (77.5, 81.0), 'region': 'North'},
    'West Bengal': {'n_districts': 23, 'lat': (21.5, 27.5), 'lon': (86.5, 89.5), 'region': 'East'},
    'Delhi': {'n_districts': 11, 'lat': (28.4, 28.9), 'lon': (76.8, 77.4), 'region': 'North'},
    'Jammu & Kashmir': {'n_districts': 20, 'lat': (32.5, 37.0), 'lon': (73.5, 80.0), 'region': 'North'},
    'Ladakh': {'n_districts': 2, 'lat': (32.5, 35.5), 'lon': (75.5, 78.5), 'region': 'North'},
    'Puducherry': {'n_districts': 4, 'lat': (10.5, 12.0), 'lon': (79.5, 80.0), 'region': 'South'},
    'Chandigarh': {'n_districts': 1, 'lat': (30.7, 30.8), 'lon': (76.7, 76.8), 'region': 'North'},
    'Andaman & Nicobar': {'n_districts': 3, 'lat': (7.0, 13.5), 'lon': (92.0, 94.0), 'region': 'Islands'},
    'Dadra & Nagar Haveli': {'n_districts': 1, 'lat': (20.1, 20.4), 'lon': (72.9, 73.3), 'region': 'West'},
    'Lakshadweep': {'n_districts': 1, 'lat': (10.0, 12.5), 'lon': (71.5, 74.0), 'region': 'Islands'},
}

def generate_district_data():
    """Generate realistic district-level multi-source data."""
    
    records = []
    district_id = 0
    
    for state, config in STATES_CONFIG.items():
        n = config['n_districts']
        region = config['region']
        
        for i in range(n):
            district_id += 1
            lat = np.random.uniform(*config['lat'])
            lon = np.random.uniform(*config['lon'])
            
            # ============================================
            # DEMOGRAPHIC FEATURES (Census 2011 distributions)
            # ============================================
            # Population (district range: 10K - 11M, log-normal)
            population = int(np.random.lognormal(mean=13.5, sigma=1.2))
            population = np.clip(population, 10000, 12000000)
            
            # Urban proportion (India avg ~31%, varies by region)
            urban_base = {'North': 0.34, 'South': 0.38, 'East': 0.18, 
                         'West': 0.42, 'Central': 0.22, 'NorthEast': 0.15, 'Islands': 0.35}
            urban_pct = np.clip(np.random.beta(2, 5) * 0.8 + urban_base.get(region, 0.3) * 0.4, 0.02, 0.99)
            
            # Population density (persons/sq.km)
            if 'Delhi' in state:
                pop_density = np.random.uniform(5000, 30000)
            elif urban_pct > 0.6:
                pop_density = np.random.lognormal(7.0, 0.8)
            else:
                pop_density = np.random.lognormal(5.5, 1.0)
            pop_density = np.clip(pop_density, 5, 35000)
            
            # Literacy rate (India avg 74%, range 50-95%)
            literacy_base = {'South': 0.82, 'North': 0.72, 'East': 0.65, 
                           'West': 0.78, 'Central': 0.68, 'NorthEast': 0.73, 'Islands': 0.86}
            literacy_rate = np.clip(
                np.random.normal(literacy_base.get(region, 0.72), 0.08), 0.40, 0.98
            )
            
            # SC/ST proportion (Census)
            sc_pct = np.clip(np.random.beta(2, 8), 0.0, 0.50)
            st_pct = np.clip(np.random.beta(1.5, 10) if region != 'NorthEast' else np.random.beta(5, 3), 0.0, 0.95)
            
            # Household size (India avg ~4.9)
            household_size = np.clip(np.random.normal(4.9, 0.8), 2.0, 9.0)
            
            # Persons per room (crowding indicator)
            persons_per_room = np.clip(np.random.lognormal(0.5, 0.4), 1.0, 6.0)
            
            # Migration rate
            migration_rate = np.clip(np.random.beta(2, 10), 0.01, 0.30)
            
            # ============================================
            # NUTRITIONAL/HEALTH (NFHS-5 distributions)
            # ============================================
            # Stunting under-5 (India avg ~35%, range 12-55%)
            stunting_rate = np.clip(
                np.random.normal(0.35 - 0.1 * (literacy_rate - 0.7), 0.08), 0.08, 0.60
            )
            
            # Anaemia prevalence (women 15-49, India avg ~57%)
            anaemia_pct = np.clip(
                np.random.normal(0.57 - 0.15 * (literacy_rate - 0.7), 0.10), 0.20, 0.85
            )
            
            # BMI underweight (<18.5, India avg ~18.7%)
            underweight_pct = np.clip(
                np.random.normal(0.187 + 0.2 * (1 - urban_pct), 0.06), 0.02, 0.45
            )
            
            # Diabetes prevalence (proxy: high blood sugar, avg ~10%)
            diabetes_pct = np.clip(
                np.random.normal(0.10 + 0.05 * urban_pct, 0.04), 0.02, 0.25
            )
            
            # HIV prevalence (India avg ~0.22%, higher in high-burden states)
            hiv_high_states = ['Maharashtra', 'Karnataka', 'Andhra Pradesh', 'Telangana', 
                              'Tamil Nadu', 'Manipur', 'Nagaland', 'Mizoram']
            hiv_base = 0.005 if state in hiv_high_states else 0.001
            hiv_prevalence = np.clip(np.random.exponential(hiv_base), 0.0001, 0.03)
            
            # BCG coverage (India avg ~92%)
            bcg_coverage = np.clip(np.random.normal(0.92, 0.06), 0.50, 0.99)
            
            # ============================================
            # SOCIOECONOMIC (SECC/Census)
            # ============================================
            # Below poverty line %
            bpl_pct = np.clip(
                np.random.normal(0.22 + 0.15 * (1 - literacy_rate), 0.08), 0.01, 0.65
            )
            
            # Female literacy gap
            female_literacy_gap = np.clip(np.random.normal(0.17, 0.06), 0.0, 0.35)
            
            # Worker participation rate
            worker_participation = np.clip(np.random.normal(0.40, 0.08), 0.15, 0.65)
            
            # Open defecation %  (India avg ~26% in 2019, varying widely)
            open_defecation = np.clip(
                np.random.beta(2, 5) * (1 - urban_pct) * 0.8, 0.0, 0.80
            )
            
            # Improved water source %
            improved_water = np.clip(
                np.random.normal(0.85 + 0.1 * urban_pct, 0.08), 0.40, 0.99
            )
            
            # Indoor cooking fuel (solid fuel %)
            solid_fuel_pct = np.clip(
                np.random.normal(0.55 - 0.3 * urban_pct, 0.12), 0.05, 0.90
            )
            
            # ============================================
            # ENVIRONMENTAL (ERA5/CPCB distributions)
            # ============================================
            # Temperature (annual mean, varies by lat)
            temp_mean = np.clip(35 - 0.5 * (lat - 8) + np.random.normal(0, 2), 15, 35)
            
            # Humidity (annual mean %)
            humidity = np.clip(
                np.random.normal(65 + 10 * (1 / (1 + np.exp(-(lon - 80)))), 10), 30, 95
            )
            
            # Rainfall (annual mm, varies by region)
            rainfall_base = {'NorthEast': 2500, 'South': 1200, 'West': 800, 
                           'North': 700, 'Central': 900, 'East': 1400, 'Islands': 3000}
            annual_rainfall = np.clip(
                np.random.lognormal(np.log(rainfall_base.get(region, 1000)), 0.4), 200, 6000
            )
            
            # PM2.5 (annual mean µg/m³)
            pm25_base = {'North': 90, 'East': 75, 'Central': 60, 'West': 45, 
                        'South': 35, 'NorthEast': 30, 'Islands': 20}
            pm25 = np.clip(
                np.random.lognormal(np.log(pm25_base.get(region, 50)), 0.4), 5, 200
            )
            
            # Altitude (meters)
            if region == 'North' and lat > 30:
                altitude = np.random.lognormal(7.5, 0.5)
            elif region == 'NorthEast':
                altitude = np.random.lognormal(6.0, 0.8)
            else:
                altitude = np.random.lognormal(5.0, 1.0)
            altitude = np.clip(altitude, 0, 6000)
            
            # ============================================
            # HEALTHCARE INFRASTRUCTURE (HMIS/NFHS)
            # ============================================
            # PHCs per 100K population
            phc_per_100k = np.clip(np.random.lognormal(1.5, 0.5), 0.5, 30)
            
            # Hospital beds per 1000
            beds_per_1000 = np.clip(
                np.random.lognormal(-0.3 + 0.5 * urban_pct, 0.5), 0.1, 10.0
            )
            
            # TB treatment facilities (DMCs) per 100K
            dmc_per_100k = np.clip(np.random.lognormal(0.5, 0.6), 0.1, 15.0)
            
            # Treatment success rate (India avg ~82%)
            treatment_success = np.clip(
                np.random.normal(0.82 + 0.1 * (beds_per_1000 / 5), 0.08), 0.50, 0.98
            )
            
            # Institutional delivery % (proxy for healthcare utilization)
            institutional_delivery = np.clip(
                np.random.normal(0.75 + 0.15 * urban_pct, 0.12), 0.20, 0.99
            )
            
            # ANC visits (4+ visits %)
            anc_4plus = np.clip(
                np.random.normal(0.58 + 0.2 * urban_pct, 0.12), 0.15, 0.95
            )
            
            # ============================================
            # TB TARGET VARIABLES (NIKSHAY/WHO-calibrated)
            # ============================================
            # TB notification rate (per 100K) - generated from risk factors
            # Key risk factors per literature: poverty, crowding, malnutrition, HIV, air quality
            tb_risk_score = (
                0.20 * (bpl_pct / 0.65) +
                0.15 * (persons_per_room / 6.0) +
                0.12 * (underweight_pct / 0.45) +
                0.12 * (stunting_rate / 0.60) +
                0.10 * (hiv_prevalence / 0.03) +
                0.08 * (pm25 / 200) +
                0.07 * (solid_fuel_pct / 0.90) +
                0.06 * (open_defecation / 0.80) +
                0.05 * (1 - treatment_success) +
                0.05 * (1 - literacy_rate)
            )
            
            # Add spatial clustering effect (nearby districts influence each other)
            spatial_noise = np.random.normal(0, 0.05)
            tb_risk_score = np.clip(tb_risk_score + spatial_noise, 0.05, 0.95)
            
            # TB notification rate (India avg ~200/100K, range 50-600)
            tb_notification_rate = np.clip(
                50 + tb_risk_score * 550 + np.random.normal(0, 30), 20, 700
            )
            
            # Case detection rate
            case_detection_rate = np.clip(
                np.random.normal(0.75 + 0.15 * (dmc_per_100k / 10), 0.08), 0.40, 0.99
            )
            
            # MDR-TB proportion (India avg ~2.8%)
            mdr_tb_pct = np.clip(np.random.exponential(0.028), 0.005, 0.15)
            
            # Hotspot label (binary: top quartile of risk)
            # Will be computed after all districts are generated
            
            records.append({
                'district_id': f'D{district_id:04d}',
                'state': state,
                'region': region,
                'latitude': round(lat, 4),
                'longitude': round(lon, 4),
                # Demographic
                'population': population,
                'urban_pct': round(urban_pct, 4),
                'pop_density': round(pop_density, 1),
                'literacy_rate': round(literacy_rate, 4),
                'sc_pct': round(sc_pct, 4),
                'st_pct': round(st_pct, 4),
                'household_size': round(household_size, 2),
                'persons_per_room': round(persons_per_room, 2),
                'migration_rate': round(migration_rate, 4),
                # Nutritional/Health
                'stunting_rate': round(stunting_rate, 4),
                'anaemia_pct': round(anaemia_pct, 4),
                'underweight_pct': round(underweight_pct, 4),
                'diabetes_pct': round(diabetes_pct, 4),
                'hiv_prevalence': round(hiv_prevalence, 6),
                'bcg_coverage': round(bcg_coverage, 4),
                # Socioeconomic
                'bpl_pct': round(bpl_pct, 4),
                'female_literacy_gap': round(female_literacy_gap, 4),
                'worker_participation': round(worker_participation, 4),
                'open_defecation': round(open_defecation, 4),
                'improved_water': round(improved_water, 4),
                'solid_fuel_pct': round(solid_fuel_pct, 4),
                # Environmental
                'temp_mean': round(temp_mean, 1),
                'humidity': round(humidity, 1),
                'annual_rainfall': round(annual_rainfall, 0),
                'pm25': round(pm25, 1),
                'altitude': round(altitude, 0),
                # Healthcare Infrastructure
                'phc_per_100k': round(phc_per_100k, 2),
                'beds_per_1000': round(beds_per_1000, 3),
                'dmc_per_100k': round(dmc_per_100k, 2),
                'treatment_success': round(treatment_success, 4),
                'institutional_delivery': round(institutional_delivery, 4),
                'anc_4plus': round(anc_4plus, 4),
                # TB Target Variables
                'tb_notification_rate': round(tb_notification_rate, 1),
                'case_detection_rate': round(case_detection_rate, 4),
                'mdr_tb_pct': round(mdr_tb_pct, 4),
                'tb_risk_score': round(tb_risk_score, 4),
            })
    
    df = pd.DataFrame(records)
    
    # Add spatial autocorrelation to TB rates (districts near each other have correlated rates)
    _add_spatial_autocorrelation(df)
    
    # Compute hotspot labels using Getis-Ord-like approach
    _compute_hotspot_labels(df)
    
    # Vulnerability index (composite)
    df['vulnerability_index'] = (
        0.3 * df['tb_risk_score_spatial'] + 
        0.2 * (df['tb_notification_rate'] / df['tb_notification_rate'].max()) +
        0.15 * (df['bpl_pct'] / df['bpl_pct'].max()) +
        0.1 * (df['persons_per_room'] / df['persons_per_room'].max()) +
        0.1 * (1 - df['treatment_success']) +
        0.1 * (df['underweight_pct'] / df['underweight_pct'].max()) +
        0.05 * (df['pm25'] / df['pm25'].max())
    )
    df['vulnerability_index'] = (df['vulnerability_index'] - df['vulnerability_index'].min()) / \
                                 (df['vulnerability_index'].max() - df['vulnerability_index'].min())
    
    return df


def _add_spatial_autocorrelation(df):
    """Add spatial autocorrelation - nearby districts should have correlated TB rates."""
    from scipy.spatial.distance import cdist
    
    coords = df[['latitude', 'longitude']].values
    dist_matrix = cdist(coords, coords, metric='euclidean')
    
    # Gaussian kernel for spatial weights (bandwidth ~2 degrees ≈ 200km)
    W = np.exp(-dist_matrix**2 / (2 * 2.0**2))
    np.fill_diagonal(W, 0)
    W = W / W.sum(axis=1, keepdims=True)
    
    # Spatially smoothed TB risk
    tb_rates = df['tb_notification_rate'].values
    spatial_smoothed = W @ tb_rates
    
    # Blend original with spatial component (Moran-style)
    alpha = 0.3  # spatial autocorrelation strength
    df['tb_notification_rate'] = np.clip(
        (1 - alpha) * tb_rates + alpha * spatial_smoothed + np.random.normal(0, 10, len(df)),
        20, 700
    )
    df['tb_risk_score_spatial'] = (df['tb_notification_rate'] - df['tb_notification_rate'].min()) / \
                                   (df['tb_notification_rate'].max() - df['tb_notification_rate'].min())


def _compute_hotspot_labels(df):
    """Compute hotspot labels using a Getis-Ord Gi*-like statistic."""
    from scipy.spatial.distance import cdist
    
    coords = df[['latitude', 'longitude']].values
    dist_matrix = cdist(coords, coords, metric='euclidean')
    
    # Binary spatial weights (within ~3 degrees ≈ 300km)
    d_threshold = 3.0
    W = (dist_matrix < d_threshold).astype(float)
    np.fill_diagonal(W, 0)
    
    tb_rates = df['tb_notification_rate'].values
    n = len(tb_rates)
    x_bar = tb_rates.mean()
    s = tb_rates.std()
    
    # Gi* statistic for each district
    gi_star = np.zeros(n)
    for i in range(n):
        wi = W[i]
        wi_sum = wi.sum()
        numerator = (wi * tb_rates).sum() - x_bar * wi_sum
        denominator = s * np.sqrt((n * (wi**2).sum() - wi_sum**2) / (n - 1))
        gi_star[i] = numerator / denominator if denominator > 0 else 0
    
    df['gi_star_zscore'] = gi_star
    
    # Hotspot: z > 1.96 (95% confidence), Coldspot: z < -1.96
    df['hotspot_label'] = 0  # Not significant
    df.loc[df['gi_star_zscore'] > 1.96, 'hotspot_label'] = 1   # Hotspot
    df.loc[df['gi_star_zscore'] < -1.96, 'hotspot_label'] = -1  # Coldspot
    
    # Binary hotspot for classification
    df['is_hotspot'] = (df['gi_star_zscore'] > 1.65).astype(int)  # More liberal threshold
    
    # Vulnerability class (4 levels)
    df['vulnerability_class'] = pd.qcut(
        df['tb_notification_rate'], q=4, labels=['Low', 'Moderate', 'High', 'Very High']
    )


if __name__ == '__main__':
    df = generate_district_data()
    print(f"Generated {len(df)} districts across {df['state'].nunique()} states")
    print(f"\nShape: {df.shape}")
    print(f"\nColumns ({len(df.columns)}):")
    for col in df.columns:
        print(f"  {col}: {df[col].dtype} | range: [{df[col].min():.4f}, {df[col].max():.4f}]" 
              if df[col].dtype in ['float64', 'int64'] else f"  {col}: {df[col].dtype}")
    
    print(f"\nHotspot distribution:")
    print(df['is_hotspot'].value_counts())
    print(f"\nVulnerability class distribution:")
    print(df['vulnerability_class'].value_counts())
    print(f"\nTB notification rate stats:")
    print(df['tb_notification_rate'].describe())
    
    df.to_csv('/app/tb_vulnerability_pipeline/district_data.csv', index=False)
    print(f"\nSaved to /app/tb_vulnerability_pipeline/district_data.csv")