"""
Spatial Feature Engineering Module
===================================
1. Geographically-aware clustering (K-Means, DBSCAN on coordinates)
2. Graph-based proximity features (k-hop neighbor aggregation)
3. Spatial lag features
4. Getis-Ord Gi* hotspot statistics
5. Spatial encoding (Fourier positional encoding for lat/lon)
"""

import numpy as np
import pandas as pd
from scipy.spatial.distance import cdist
from sklearn.cluster import KMeans, DBSCAN
import networkx as nx


def build_adjacency_graph(df, method='knn', k=6, distance_threshold=2.5):
    """Build spatial adjacency graph from district coordinates.
    
    Args:
        method: 'knn' (k-nearest neighbors) or 'distance' (distance threshold)
        k: number of neighbors for knn
        distance_threshold: max distance in degrees for distance-based graph
    """
    coords = df[['latitude', 'longitude']].values
    dist_matrix = cdist(coords, coords, metric='euclidean')
    
    G = nx.Graph()
    n = len(df)
    
    for i in range(n):
        G.add_node(i, **{
            'district_id': df.iloc[i]['district_id'],
            'state': df.iloc[i]['state'],
            'lat': df.iloc[i]['latitude'],
            'lon': df.iloc[i]['longitude']
        })
    
    if method == 'knn':
        for i in range(n):
            neighbors = np.argsort(dist_matrix[i])[1:k+1]
            for j in neighbors:
                G.add_edge(i, j, weight=1.0 / (1.0 + dist_matrix[i, j]))
    elif method == 'distance':
        for i in range(n):
            for j in range(i+1, n):
                if dist_matrix[i, j] < distance_threshold:
                    G.add_edge(i, j, weight=1.0 / (1.0 + dist_matrix[i, j]))
    
    return G, dist_matrix


def compute_spatial_lag_features(df, dist_matrix, feature_cols, bandwidth=2.0):
    """Compute spatial lag features: weighted mean of neighbors' values.
    
    Uses Gaussian kernel weighting.
    """
    coords = df[['latitude', 'longitude']].values
    
    # Gaussian kernel weights
    W = np.exp(-dist_matrix**2 / (2 * bandwidth**2))
    np.fill_diagonal(W, 0)
    W_norm = W / W.sum(axis=1, keepdims=True)
    
    lag_features = {}
    for col in feature_cols:
        values = df[col].values.astype(float)
        lag_features[f'{col}_spatial_lag'] = W_norm @ values
        
        # Spatial variance (local heterogeneity)
        mean_neighbor = W_norm @ values
        sq_diff = W_norm @ (values**2) - mean_neighbor**2
        lag_features[f'{col}_spatial_var'] = np.maximum(sq_diff, 0)
    
    return pd.DataFrame(lag_features, index=df.index)


def compute_graph_proximity_features(G, df, feature_cols, k_hops=[1, 2, 3]):
    """Compute k-hop aggregation features on the spatial graph.
    
    For each node, aggregate feature values from 1-hop, 2-hop, 3-hop neighbors.
    """
    n = len(df)
    features = {}
    
    for k in k_hops:
        for col in feature_cols:
            values = df[col].values.astype(float)
            mean_vals = np.zeros(n)
            max_vals = np.zeros(n)
            min_vals = np.zeros(n)
            std_vals = np.zeros(n)
            count_vals = np.zeros(n)
            
            for node in G.nodes():
                # Get all nodes within k hops
                neighbors = set()
                for target, length in nx.single_source_shortest_path_length(G, node, cutoff=k).items():
                    if target != node and length <= k:
                        neighbors.add(target)
                
                if neighbors:
                    neighbor_vals = values[list(neighbors)]
                    mean_vals[node] = neighbor_vals.mean()
                    max_vals[node] = neighbor_vals.max()
                    min_vals[node] = neighbor_vals.min()
                    std_vals[node] = neighbor_vals.std() if len(neighbor_vals) > 1 else 0
                    count_vals[node] = len(neighbors)
                else:
                    mean_vals[node] = values[node]
                    max_vals[node] = values[node]
                    min_vals[node] = values[node]
            
            features[f'{col}_{k}hop_mean'] = mean_vals
            features[f'{col}_{k}hop_max'] = max_vals
            features[f'{col}_{k}hop_range'] = max_vals - min_vals
        
        # Add neighbor count at each hop level
        features[f'neighbor_count_{k}hop'] = count_vals
    
    return pd.DataFrame(features, index=df.index)


def geographic_clustering(df, methods=['kmeans', 'dbscan'], 
                          kmeans_clusters=[5, 10, 20, 50],
                          dbscan_eps=[1.5, 2.5], dbscan_min_samples=3):
    """Apply multiple geographic clustering methods.
    
    Returns cluster labels for each method/config.
    """
    coords = df[['latitude', 'longitude']].values
    cluster_features = {}
    
    # K-Means clustering at different granularities
    for n_clusters in kmeans_clusters:
        km = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
        labels = km.fit_predict(coords)
        cluster_features[f'geo_cluster_kmeans_{n_clusters}'] = labels
        
        # Distance to cluster center
        centers = km.cluster_centers_
        dists = np.array([np.linalg.norm(coords[i] - centers[labels[i]]) 
                         for i in range(len(df))])
        cluster_features[f'dist_to_center_kmeans_{n_clusters}'] = dists
    
    # DBSCAN (density-based)
    for eps in dbscan_eps:
        db = DBSCAN(eps=eps, min_samples=dbscan_min_samples)
        labels = db.fit_predict(coords)
        cluster_features[f'geo_cluster_dbscan_eps{eps}'] = labels
        n_clusters = len(set(labels)) - (1 if -1 in labels else 0)
        print(f"  DBSCAN eps={eps}: {n_clusters} clusters, {(labels == -1).sum()} outliers")
    
    return pd.DataFrame(cluster_features, index=df.index)


def fourier_positional_encoding(df, n_frequencies=16):
    """Fourier positional encoding for latitude/longitude.
    
    Maps (lat, lon) → high-dimensional sinusoidal features.
    Similar to NeRF positional encoding.
    """
    lat = df['latitude'].values
    lon = df['longitude'].values
    
    # Normalize to [0, 1]
    lat_norm = (lat - lat.min()) / (lat.max() - lat.min())
    lon_norm = (lon - lon.min()) / (lon.max() - lon.min())
    
    features = {}
    for i in range(n_frequencies):
        freq = 2.0 ** i
        features[f'lat_sin_{i}'] = np.sin(2 * np.pi * freq * lat_norm)
        features[f'lat_cos_{i}'] = np.cos(2 * np.pi * freq * lat_norm)
        features[f'lon_sin_{i}'] = np.sin(2 * np.pi * freq * lon_norm)
        features[f'lon_cos_{i}'] = np.cos(2 * np.pi * freq * lon_norm)
    
    return pd.DataFrame(features, index=df.index)


def compute_spatial_autocorrelation_features(df, dist_matrix, 
                                              target_col='tb_notification_rate',
                                              bandwidths=[1.0, 2.0, 3.0, 5.0]):
    """Multi-scale spatial autocorrelation features.
    
    Computes local Moran's I-like statistics at multiple bandwidths.
    """
    values = df[target_col].values
    z = (values - values.mean()) / values.std()
    n = len(values)
    
    features = {}
    for bw in bandwidths:
        W = np.exp(-dist_matrix**2 / (2 * bw**2))
        np.fill_diagonal(W, 0)
        W_norm = W / W.sum(axis=1, keepdims=True)
        
        # Local Moran's I: z_i * sum(w_ij * z_j)
        spatial_lag_z = W_norm @ z
        local_moran = z * spatial_lag_z
        
        features[f'local_moran_bw{bw}'] = local_moran
        features[f'spatial_lag_z_bw{bw}'] = spatial_lag_z
        
        # LISA quadrant (HH=1, LH=2, LL=3, HL=4)
        quadrant = np.zeros(n, dtype=int)
        quadrant[(z > 0) & (spatial_lag_z > 0)] = 1  # High-High
        quadrant[(z < 0) & (spatial_lag_z > 0)] = 2  # Low-High
        quadrant[(z < 0) & (spatial_lag_z < 0)] = 3  # Low-Low
        quadrant[(z > 0) & (spatial_lag_z < 0)] = 4  # High-Low
        features[f'lisa_quadrant_bw{bw}'] = quadrant
    
    return pd.DataFrame(features, index=df.index)


def compute_all_spatial_features(df, target_col='tb_notification_rate'):
    """Master function to compute all spatial features."""
    
    print("Building spatial adjacency graph...")
    G, dist_matrix = build_adjacency_graph(df, method='knn', k=6)
    print(f"  Graph: {G.number_of_nodes()} nodes, {G.number_of_edges()} edges")
    
    # Select key features for spatial lag computation
    key_features = [
        'tb_notification_rate', 'bpl_pct', 'persons_per_room',
        'underweight_pct', 'pm25', 'literacy_rate', 'urban_pct',
        'treatment_success', 'pop_density', 'hiv_prevalence'
    ]
    
    print("Computing spatial lag features...")
    lag_df = compute_spatial_lag_features(df, dist_matrix, key_features, bandwidth=2.0)
    
    print("Computing graph proximity features...")
    graph_key = ['tb_notification_rate', 'bpl_pct', 'underweight_pct', 'pm25', 'pop_density']
    graph_df = compute_graph_proximity_features(G, df, graph_key, k_hops=[1, 2])
    
    print("Computing geographic clusters...")
    cluster_df = geographic_clustering(df)
    
    print("Computing Fourier positional encoding...")
    fourier_df = fourier_positional_encoding(df, n_frequencies=8)
    
    print("Computing spatial autocorrelation features...")
    autocorr_df = compute_spatial_autocorrelation_features(
        df, dist_matrix, target_col=target_col
    )
    
    # Combine all spatial features
    spatial_df = pd.concat([lag_df, graph_df, cluster_df, fourier_df, autocorr_df], axis=1)
    
    print(f"Total spatial features generated: {spatial_df.shape[1]}")
    return spatial_df, G, dist_matrix


if __name__ == '__main__':
    df = pd.read_csv('/app/tb_vulnerability_pipeline/district_data.csv')
    spatial_df, G, dist_matrix = compute_all_spatial_features(df)
    
    print(f"\nSpatial features shape: {spatial_df.shape}")
    print(f"Sample features:\n{spatial_df.head()}")
    
    # Save
    full_df = pd.concat([df, spatial_df], axis=1)
    full_df.to_csv('/app/tb_vulnerability_pipeline/district_data_with_spatial.csv', index=False)
    print(f"\nSaved full dataset with {full_df.shape[1]} columns")