import os
import tempfile
import pandas as pd
import numpy as np
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler
import folium
from folium.plugins import HeatMap
import gradio as gr

def haversine_distance(lat1, lon1, lat2, lon2):
    """Calculates geodesic distance in kilometers."""
    lat1, lon1, lat2, lon2 = map(np.radians, [lat1, lon1, lat2, lon2])
    dlat = lat2 - lat1
    dlon = lon2 - lon1
    a = np.sin(dlat/2.0)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon/2.0)**2
    c = 2.0 * np.arcsin(np.sqrt(a))
    return c * 6371.0

def run_idw_interpolation(sampled_coords, sampled_vals, target_coords, power=2.0):
    """Estimates values at target coordinates using Inverse Distance Weighting (IDW)."""
    n_targets = len(target_coords)
    n_samples = len(sampled_coords)
    
    estimates = []
    for i in range(n_targets):
        t_lat, t_lon = target_coords[i]
        
        # Calculate distances to all sampled points
        dists = haversine_distance(t_lat, t_lon, sampled_coords[:, 0], sampled_coords[:, 1])
        
        # Handle exact overlap (distance = 0)
        zero_idx = np.where(dists == 0.0)[0]
        if len(zero_idx) > 0:
            estimates.append(sampled_vals[zero_idx[0]])
            continue
            
        # Compute weights: w_i = 1 / (d_i^p)
        weights = 1.0 / (dists ** power)
        
        # Estimate value
        est_val = np.sum(weights * sampled_vals) / np.sum(weights)
        estimates.append(est_val)
        
    return np.array(estimates)

def run_geodemographic_clustering(df, cols, num_clusters=4):
    """Partitions tracts into geodemographic segments using K-Means."""
    try:
        df_sub = df[cols].dropna()
        if len(df_sub) < num_clusters:
            return None, "Error: Dataset must contain at least as many rows as clusters."
            
        # Normalize features
        scaler = StandardScaler()
        scaled_features = scaler.fit_transform(df_sub)
        
        # Run K-Means
        kmeans = KMeans(n_clusters=num_clusters, random_state=42, n_init=10)
        labels = kmeans.fit_predict(scaled_features)
        
        df_clustered = df.dropna(subset=cols).copy()
        df_clustered["Cluster_ID"] = labels
        # Convert to readable cluster labels
        df_clustered["Cluster_Name"] = [f"Segment {l + 1}" for l in labels]
        
        # Compile cluster profile centroids
        centroids = scaler.inverse_transform(kmeans.cluster_centers_)
        df_profiles = pd.DataFrame(centroids, columns=cols)
        df_profiles.insert(0, "Cluster Segment", [f"Segment {i + 1}" for i in range(num_clusters)])
        
        return df_profiles, df_clustered
    except Exception as e:
        print(f"Clustering failed: {e}")
        return None, f"Clustering processing failed: {e}"

def generate_interpolation_surface_map(df_sampled, lat_col, lon_col, val_col, df_estimated, power):
    """Draws a beautiful color-graded Folium map with original values and the estimated grid."""
    mean_lat = df_sampled[lat_col].mean()
    mean_lon = df_sampled[lon_col].mean()
    
    m = folium.Map(location=[mean_lat, mean_lon], zoom_start=11, tiles="CartoDB dark_matter")
    
    # 1. Overlay estimated points as a translucent coordinate surface
    # Min/max values for gradient color mapping
    v_min = min(df_sampled[val_col].min(), df_estimated["Estimated_Value"].min())
    v_max = max(df_sampled[val_col].max(), df_estimated["Estimated_Value"].max())
    v_range = max(v_max - v_min, 0.0001)
    
    for i, row in df_estimated.iterrows():
        lat = row["Latitude"]
        lon = row["Longitude"]
        val = row["Estimated_Value"]
        
        # Color gradient calculation (Spectral/Jet style: Blue (Low) to Red (High))
        ratio = (val - v_min) / v_range
        # Map ratio to a simple rgb color
        r = int(255 * ratio)
        g = int(255 * (1.0 - abs(ratio - 0.5) * 2.0))
        b = int(255 * (1.0 - ratio))
        color_hex = f"#{r:02x}{g:02x}{b:02x}"
        
        folium.CircleMarker(
            location=[lat, lon],
            radius=4,
            color=color_hex,
            fill=True,
            fill_color=color_hex,
            fill_opacity=0.4,
            weight=0,
            popup=f"Estimated: {val:.2f}"
        ).add_to(m)
        
    # 2. Overlay original measured sampled markers (larger, high contrast)
    for i, row in df_sampled.iterrows():
        lat = row[lat_col]
        lon = row[lon_col]
        val = row[val_col]
        
        folium.CircleMarker(
            location=[lat, lon],
            radius=9,
            color="#ffffff", # White border
            fill=True,
            fill_color="#eab308", # Amber fill
            fill_opacity=0.9,
            weight=2,
            popup=f"<b>Measured Sample</b><br>Value: {val:.2f}"
        ).add_to(m)
        
    return m

def generate_geodemographic_map(df_clustered, lat_col, lon_col, num_clusters):
    """Draws a Folium map showing color-coded geodemographic segments."""
    mean_lat = df_clustered[lat_col].mean()
    mean_lon = df_clustered[lon_col].mean()
    
    m = folium.Map(location=[mean_lat, mean_lon], zoom_start=11, tiles="CartoDB positron")
    
    # Palette of distinct segment colors
    colors = ["#10b981", "#3b82f6", "#ef4444", "#8b5cf6", "#f59e0b", "#ec4899", "#14b8a6", "#6b7280"]
    
    for i, row in df_clustered.iterrows():
        lat = row[lat_col]
        lon = row[lon_col]
        segment = row["Cluster_Name"]
        c_id = int(row["Cluster_ID"])
        color = colors[c_id % len(colors)]
        
        popup_html = f"""
        <div style="font-family: 'Inter', sans-serif; color: #111827; min-width: 150px;">
            <h4 style="margin:0 0 5px 0;">Geodemographic Unit</h4>
            <b>Assigned Class</b>: <span style="color: {color}; font-weight: bold;">{segment}</span>
        </div>
        """
        
        folium.CircleMarker(
            location=[lat, lon],
            radius=8,
            color=color,
            fill=True,
            fill_color=color,
            fill_opacity=0.7,
            weight=1.5,
            popup=folium.Popup(popup_html, max_width=200)
        ).add_to(m)
        
    return m

def full_pipeline_interpolation(file_sampled, file_targets, lat_col, lon_col, val_col, power):
    """Main pipeline execution for spatial interpolation."""
    if file_sampled is None:
        return None, "Please upload the Sampled Data CSV file.", pd.DataFrame(), None
        
    try:
        df_sampled = pd.read_csv(file_sampled.name)
        
        # Column checks
        for c in [lat_col, lon_col, val_col]:
            if c not in df_sampled.columns:
                return None, f"ERROR: Column '{c}' not found in sampled CSV!", pd.DataFrame(), None
                
        df_sampled_clean = df_sampled.dropna(subset=[lat_col, lon_col, val_col]).copy()
        
        # Set target coordinates
        if file_targets is not None:
            df_targets = pd.read_csv(file_targets.name)
            # Check columns
            if lat_col not in df_targets.columns or lon_col not in df_targets.columns:
                return None, f"ERROR: Latitude/Longitude columns not found in target CSV!", pd.DataFrame(), None
            df_targets_clean = df_targets.dropna(subset=[lat_col, lon_col]).copy()
            target_coords = df_targets_clean[[lat_col, lon_col]].values
            df_estimated = pd.DataFrame({
                "Latitude": target_coords[:, 0],
                "Longitude": target_coords[:, 1]
            })
        else:
            # Auto-generate a dense bounding-box coordinate estimation grid (15x15 surface)
            min_lat = df_sampled_clean[lat_col].min()
            max_lat = df_sampled_clean[lat_col].max()
            min_lon = df_sampled_clean[lon_col].min()
            max_lon = df_sampled_clean[lon_col].max()
            
            # Margin buffers
            margin_lat = max(0.005, (max_lat - min_lat) * 0.05)
            margin_lon = max(0.005, (max_lon - min_lon) * 0.05)
            
            lat_grid = np.linspace(min_lat - margin_lat, max_lat + margin_lat, 15)
            lon_grid = np.linspace(min_lon - margin_lon, max_lon + margin_lon, 15)
            
            grid_lats, grid_lons = np.meshgrid(lat_grid, lon_grid)
            target_coords = np.column_stack((grid_lats.ravel(), grid_lons.ravel()))
            df_estimated = pd.DataFrame({
                "Latitude": target_coords[:, 0],
                "Longitude": target_coords[:, 1]
            })
            
        sampled_coords = df_sampled_clean[[lat_col, lon_col]].values
        sampled_vals = df_sampled_clean[val_col].values
        
        # Calculate IDW
        estimated_values = run_idw_interpolation(sampled_coords, sampled_vals, target_coords, power)
        df_estimated["Estimated_Value"] = estimated_values
        
        # Draw map
        map_obj = generate_interpolation_surface_map(
            df_sampled_clean, lat_col, lon_col, val_col, df_estimated, power
        )
        
        # Save HTML map
        temp_map = tempfile.NamedTemporaryFile(delete=False, suffix=".html")
        map_obj.save(temp_map.name)
        
        status_md = f"""
        ### 📊 Spatial Interpolation (IDW) Complete:
        * **Power Parameter ($p$)**: `{power:.1f}`
        * **Sampled Points**: `{len(df_sampled_clean)}`
        * **Interpolated Target Grid Points**: `{len(df_estimated)}`
        * **Estimated Value Range**: `{estimated_values.min():.2f}` to `{estimated_values.max():.2f}`
        
        *Interpretation*: The Leaflet map displays estimated values as a continuous spectral gradient surface. Standard IDW power ($p=2.0$) represents typical decay. Try adjusting the slider to observe localized vs. global smoothing.
        """
        
        # CSV download path
        temp_csv = tempfile.NamedTemporaryFile(delete=False, suffix="_interpolated_data.csv")
        df_estimated.to_csv(temp_csv.name, index=False)
        
        return temp_map.name, status_md, df_estimated, temp_csv.name
    except Exception as e:
        return None, f"Interpolation processing failed: {e}", pd.DataFrame(), None

def full_pipeline_clustering(file_sampled, lat_col, lon_col, cluster_cols_str, num_clusters):
    """Main pipeline execution for geodemographic clustering."""
    if file_sampled is None:
        return None, "Please upload the Census Demographic CSV file.", pd.DataFrame(), None
        
    try:
        df = pd.read_csv(file_sampled.name)
        
        # Parse columns
        cols = [c.strip() for c in cluster_cols_str.split(",") if c.strip() in df.columns]
        if not cols:
            return None, "ERROR: None of the entered attributes were found in the uploaded CSV! Check column names.", pd.DataFrame(), None
            
        df_profiles, df_clustered = run_geodemographic_clustering(df, cols, num_clusters)
        
        if df_profiles is None:
            # df_profiles holds error string
            return None, df_clustered, pd.DataFrame(), None
            
        # Draw map
        map_obj = generate_geodemographic_map(df_clustered, lat_col, lon_col, num_clusters)
        
        # Save HTML map
        temp_map = tempfile.NamedTemporaryFile(delete=False, suffix=".html")
        map_obj.save(temp_map.name)
        
        status_md = f"""
        ### 📊 Geodemographic Segment Clustering Complete:
        * **Total Features Clustered**: `{len(df_clustered)}`
        * **Target Segments (k)**: `{num_clusters}`
        * **Analyzed Attributes**: `{', '.join(cols)}`
        
        *Interpretation*: View the map to see the geographical spatial distribution of neighborhood types. The table below outlines the average profile characteristics (centroids) of each segment!
        """
        
        # CSV download path
        temp_csv = tempfile.NamedTemporaryFile(delete=False, suffix="_geodemographics.csv")
        df_clustered.to_csv(temp_csv.name, index=False)
        
        return temp_map.name, status_md, df_profiles, temp_csv.name
    except Exception as e:
        return None, f"Clustering processing failed: {e}", pd.DataFrame(), None

# Premium Monochrome / custom Yellow styling
custom_css = """
body { background-color: #0d0f12; color: #e3e6eb; font-family: 'Inter', sans-serif; }
.gradio-container { max-width: 1200px !important; margin: 0 auto !important; }
h1, h2, h3 { color: #ffffff !important; font-weight: 700 !important; }
.btn-primary { background: linear-gradient(135deg, #eab308 0%, #ca8a04 100%) !important; border: none !important; color: black !important; font-weight: 600 !important; }
.btn-primary:hover { filter: brightness(1.1); }
"""

with gr.Blocks(theme=gr.themes.Monochrome(), css=custom_css) as demo:
    gr.Markdown(
        """
        # 🔆 Spatial Interpolation & Geodemographic Modeler
        ### Estimate values in un-sampled locations based on sampled statistics using Inverse Distance Weighting (IDW), or partition geographic areas into geodemographic segments using K-Means clustering.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=4):
            with gr.Card():
                gr.Markdown("### 1. Load Primary Spatial Dataset")
                file_sampled_input = gr.File(label="Upload Sampled/Census CSV", file_types=[".csv"])
                with gr.Row():
                    lat_name = gr.Textbox(label="Latitude Column Name", value="Latitude")
                    lon_name = gr.Textbox(label="Longitude Column Name", value="Longitude")
                    
            with gr.Tabs():
                with gr.TabItem("🔆 IDW Spatial Interpolation"):
                    val_name = gr.Textbox(label="Target Variable to Interpolate", value="AQI")
                    file_targets_input = gr.File(label="Upload Un-sampled Target coordinates (Optional, leaves blank to auto-grid)", file_types=[".csv"])
                    idw_power = gr.Slider(
                        minimum=1.0, maximum=3.0, value=2.0, step=0.1,
                        label="IDW Distance Power Parameter (p)"
                    )
                    interpolate_btn = gr.Button("Generate Estimated Surface Map", variant="primary", elem_classes="btn-primary")
                    
                with gr.TabItem("📊 Geodemographic Clustering"):
                    cluster_cols = gr.Textbox(
                        label="Features to Cluster (Comma-separated column names)",
                        value="PovertyRate, UnemploymentRate, MedianRent"
                    )
                    num_classes = gr.Slider(
                        minimum=2, maximum=8, value=4, step=1,
                        label="Target Geodemographic Segments (k)"
                    )
                    cluster_btn = gr.Button("Generate Neighborhood Clusters", variant="primary", elem_classes="btn-primary")
                    
        with gr.Column(scale=6):
            with gr.Tabs():
                with gr.TabItem("🗺️ Interactive Leaflet Modeler Map"):
                    map_output = gr.HTML(label="Leaflet Interpolation/Cluster Grid", value="<div style='text-align: center; padding: 50px; color: gray;'>Map will load here...</div>")
                    summary_output = gr.Markdown("Please load datasets and run spatial modeling above.")
                    
                with gr.TabItem("📊 Analytical Modeling Reports"):
                    table_output = gr.Dataframe(
                        label="Interpolated Estimations Grid / Centroids Table",
                        interactive=False,
                        wrap=True
                    )
                    download_btn = gr.File(label="Download Labeled CSV Database", interactive=False)

    # Interpolation trigger
    interpolate_btn.click(
        fn=full_pipeline_interpolation,
        inputs=[file_sampled_input, file_targets_input, lat_name, lon_name, val_name, idw_power],
        outputs=[map_output, summary_output, table_output, download_btn]
    )
    
    # Clustering trigger
    cluster_btn.click(
        fn=full_pipeline_clustering,
        inputs=[file_sampled_input, lat_name, lon_name, cluster_cols, num_classes],
        outputs=[map_output, summary_output, table_output, download_btn]
    )

if __name__ == "__main__":
    demo.launch()