import os
import tempfile
import pandas as pd
import numpy as np
import folium
import gradio as gr

def haversine_distance(lat1, lon1, lat2, lon2):
    """Calculates the great-circle distance between two points on the Earth in kilometers."""
    # Convert decimal degrees to radians
    lat1, lon1, lat2, lon2 = map(np.radians, [lat1, lon1, lat2, lon2])
    
    # Haversine formula
    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))
    r = 6371.0 # Radius of Earth in kilometers
    return c * r

def run_proximity_audit(df_dem, df_poi, dem_lat, dem_lon, poi_lat, poi_lon, radius_km):
    """Audits demographic attributes inside vs outside the point-of-interest buffers."""
    try:
        n_dem = len(df_dem)
        n_poi = len(df_poi)
        
        if n_dem == 0 or n_poi == 0:
            return None, "Error: Both datasets must contain at least 1 record."
            
        # Extract coordinates
        dem_lats = df_dem[dem_lat].values
        dem_lons = df_dem[dem_lon].values
        
        poi_lats = df_poi[poi_lat].values
        poi_lons = df_poi[poi_lon].values
        
        # Calculate minimum distance from each demographic point to any POI
        min_distances = []
        for i in range(n_dem):
            # Compute distance to all POIs
            dists = haversine_distance(dem_lats[i], dem_lons[i], poi_lats, poi_lons)
            min_distances.append(np.min(dists))
            
        min_distances = np.array(min_distances)
        
        # Categorize
        inside_buffer = min_distances <= radius_km
        
        df_audit = df_dem.copy()
        df_audit["Distance_to_POI_km"] = min_distances
        df_audit["Location_Class"] = np.where(inside_buffer, "Inside Buffer", "Outside Buffer")
        
        # Compile demographic comparisons
        # Identify all numerical columns excluding lat, lon
        exclude_cols = [dem_lat, dem_lon, "Distance_to_POI_km", "Location_Class"]
        num_cols = [c for c in df_dem.columns if pd.api.types.is_numeric_dtype(df_dem[c]) and c not in exclude_cols]
        
        comparisons = []
        for col in num_cols:
            mean_inside = df_audit[df_audit["Location_Class"] == "Inside Buffer"][col].mean()
            mean_outside = df_audit[df_audit["Location_Class"] == "Outside Buffer"][col].mean()
            
            # Handle empty divisions or NaNs
            mean_inside = 0.0 if np.isnan(mean_inside) else mean_inside
            mean_outside = 0.0 if np.isnan(mean_outside) else mean_outside
            
            # Disparity percentage
            if mean_outside != 0:
                disparity_pct = ((mean_inside - mean_outside) / mean_outside) * 100.0
            else:
                disparity_pct = 0.0
                
            comparisons.append({
                "Demographic Attribute": col,
                "Average (Inside Buffer)": mean_inside,
                "Average (Outside Buffer)": mean_outside,
                "Relative Disparity (%)": disparity_pct
            })
            
        df_compare = pd.DataFrame(comparisons)
        return df_compare, df_audit
    except Exception as e:
        print(f"Audit error: {e}")
        return None, f"Audit processing failed: {e}"

def generate_proximity_map(df_audit, df_poi, dem_lat, dem_lon, poi_lat, poi_lon, poi_label, radius_km):
    """Draws a beautiful Folium map with transparent circular buffers around POIs."""
    mean_lat = df_poi[poi_lat].mean()
    mean_lon = df_poi[poi_lon].mean()
    
    m = folium.Map(location=[mean_lat, mean_lon], zoom_start=12, tiles="CartoDB positron")
    
    # 1. Plot buffers and POI markers
    for i, row in df_poi.iterrows():
        lat = row[poi_lat]
        lon = row[poi_lon]
        label = row[poi_label] if poi_label in df_poi.columns else "Point of Interest"
        
        # Add transparent buffer overlay (radius in meters)
        folium.Circle(
            location=[lat, lon],
            radius=radius_km * 1000.0,
            color="#ef4444",
            fill=True,
            fill_color="#fca5a5",
            fill_opacity=0.2,
            weight=1
        ).add_to(m)
        
        # Add POI marker
        folium.Marker(
            location=[lat, lon],
            popup=f"<b>{label}</b>",
            icon=folium.Icon(color="red", icon="info-sign")
        ).add_to(m)
        
    # 2. Plot demographic centroids
    for i, row in df_audit.iterrows():
        lat = row[dem_lat]
        lon = row[dem_lon]
        loc_class = row["Location_Class"]
        dist = row["Distance_to_POI_km"]
        
        color = "#10b981" if loc_class == "Inside Buffer" else "#6b7280" # Green inside, Gray outside
        
        popup_html = f"""
        <div style="font-family: 'Inter', sans-serif; color: #111827; min-width: 150px;">
            <h4 style="margin:0 0 5px 0;">Centroid Area</h4>
            <b>Status</b>: {loc_class}<br>
            <b>Distance</b>: {dist:.2f} km
        </div>
        """
        
        folium.CircleMarker(
            location=[lat, lon],
            radius=6,
            color=color,
            fill=True,
            fill_color=color,
            fill_opacity=0.6,
            weight=1,
            popup=folium.Popup(popup_html, max_width=200)
        ).add_to(m)
        
    return m

def full_proximity_pipeline(file_dem, file_poi, dem_lat, dem_lon, poi_lat, poi_lon, poi_label, radius_km):
    """Executes the full loading, auditing, mapping, and download setup."""
    if file_dem is None or file_poi is None:
        return None, "Please upload both the Demographic CSV and Point of Interest CSV files.", pd.DataFrame(), None
        
    try:
        df_dem = pd.read_csv(file_dem.name)
        df_poi = pd.read_csv(file_poi.name)
        
        # Column checks
        for c in [dem_lat, dem_lon]:
            if c not in df_dem.columns:
                return None, f"ERROR: Demographic column '{c}' not found! Check columns.", pd.DataFrame(), None
        for c in [poi_lat, poi_lon]:
            if c not in df_poi.columns:
                return None, f"ERROR: POI column '{c}' not found! Check columns.", pd.DataFrame(), None
                
        df_dem_clean = df_dem.dropna(subset=[dem_lat, dem_lon]).copy()
        df_poi_clean = df_poi.dropna(subset=[poi_lat, poi_lon]).copy()
        
        df_compare, df_audit = run_proximity_audit(
            df_dem_clean, df_poi_clean, dem_lat, dem_lon, poi_lat, poi_lon, radius_km
        )
        
        if df_compare is None:
            # df_compare holds the error string
            return None, df_audit, pd.DataFrame(), None
            
        # Draw map
        map_obj = generate_proximity_map(
            df_audit, df_poi_clean, dem_lat, dem_lon, poi_lat, poi_lon, poi_label, radius_km
        )
        
        # Save HTML map
        temp_map = tempfile.NamedTemporaryFile(delete=False, suffix=".html")
        map_obj.save(temp_map.name)
        
        inside_count = len(df_audit[df_audit["Location_Class"] == "Inside Buffer"])
        total_count = len(df_audit)
        
        status_md = f"""
        ### 📊 Proximity Audit Metrics:
        * **Target POI Radius**: `{radius_km:.2f} km`
        * **Total Census Tracts/Centroids**: `{total_count}`
        * **Tracts Inside Buffer Zone**: `{inside_count} ({inside_count/total_count:.1%})`
        * **Tracts Outside Buffer Zone**: `{total_count - inside_count} ({(total_count - inside_count)/total_count:.1%})`
        
        *Interpretation*: The table on the right displays the average demographics of neighborhoods located within vs outside this buffer. Look at **Relative Disparity (%)** to detect unequal exposure or access gaps!
        """
        
        # Create CSV download path
        temp_csv = tempfile.NamedTemporaryFile(delete=False, suffix="_equity_audit.csv")
        df_audit.to_csv(temp_csv.name, index=False)
        
        return temp_map.name, status_md, df_compare, temp_csv.name
    except Exception as e:
        return None, f"Audit processing failed: {e}", pd.DataFrame(), None

# Premium Monochrome / custom Green 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, #3b82f6 0%, #1d4ed8 100%) !important; border: none !important; color: white !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(
        """
        # 📏 Proximity & Buffer Analyzer
        ### Audit spatial equity, food deserts, or environmental exposure. Draw radial distance buffers around assets or hazards, and automatically compare demographic attributes inside vs. outside the zones.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=4):
            with gr.Card():
                gr.Markdown("### 1. Upload Socio-Demographic Regions (Tracts)")
                file_dem_input = gr.File(label="Upload CSV with Latitude, Longitude, and demographics", file_types=[".csv"])
                with gr.Row():
                    dem_lat_name = gr.Textbox(label="Demographic Lat Column", value="Latitude")
                    dem_lon_name = gr.Textbox(label="Demographic Lon Column", value="Longitude")
                    
            with gr.Card():
                gr.Markdown("### 2. Upload Points of Interest (Assets/Hazards)")
                file_poi_input = gr.File(label="Upload POI CSV", file_types=[".csv"])
                with gr.Row():
                    poi_lat_name = gr.Textbox(label="POI Lat Column", value="Latitude")
                    poi_lon_name = gr.Textbox(label="POI Lon Column", value="Longitude")
                poi_lbl_name = gr.Textbox(label="POI Label/Name Column", value="Name")
                
            with gr.Card():
                gr.Markdown("### 3. Buffer Parameters")
                radius_slider = gr.Slider(
                    minimum=0.1, maximum=20.0, value=2.0, step=0.1,
                    label="Radial Proximity Buffer (km)"
                )
                analyze_btn = gr.Button("Calculate Proximity Disparities", variant="primary", elem_classes="btn-primary")
                
        with gr.Column(scale=6):
            with gr.Tabs():
                with gr.TabItem("🗺️ Proximity Leaflet Map"):
                    map_output = gr.HTML(label="Leaflet Map Grid", value="<div style='text-align: center; padding: 50px; color: gray;'>Map will load here...</div>")
                    summary_output = gr.Markdown("Please upload data and run proximity audit.")
                    
                with gr.TabItem("📊 Comparative Equity Report"):
                    table_output = gr.Dataframe(
                        label="Calculated Disparities Table (Inside vs. Outside Buffer)",
                        interactive=False,
                        wrap=True
                    )
                    download_btn = gr.File(label="Download Labeled CSV Database", interactive=False)

    analyze_btn.click(
        fn=full_proximity_pipeline,
        inputs=[file_dem_input, file_poi_input, dem_lat_name, dem_lon_name, poi_lat_name, poi_lon_name, poi_lbl_name, radius_slider],
        outputs=[map_output, summary_output, table_output, download_btn]
    )

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