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 great-circle 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_network_equity_audit(df_start, df_dest, start_lat, start_lon, dest_lat, dest_lon, dest_label, circuity_factor=1.3):
    """Calculates route distance to closest destinations and correlates with demographics."""
    try:
        n_start = len(df_start)
        n_dest = len(df_dest)
        
        if n_start == 0 or n_dest == 0:
            return None, "Error: Both datasets must contain at least 1 record."
            
        start_lats = df_start[start_lat].values
        start_lons = df_start[start_lon].values
        
        dest_lats = df_dest[dest_lat].values
        dest_lons = df_dest[dest_lon].values
        
        closest_dest_idx = []
        simulated_distances = []
        
        for i in range(n_start):
            # Calculate geodesic distance to all destinations
            dists = haversine_distance(start_lats[i], start_lons[i], dest_lats, dest_lons)
            
            # Apply urban circuity factor (standard road network winding multiplier)
            dists_adjusted = dists * circuity_factor
            
            min_idx = np.argmin(dists_adjusted)
            closest_dest_idx.append(min_idx)
            simulated_distances.append(dists_adjusted[min_idx])
            
        df_audit = df_start.copy()
        df_audit["Nearest_Destination_Index"] = closest_dest_idx
        # Destination names
        dest_names = df_dest[dest_label].values if dest_label in df_dest.columns else [f"Hub_{j}" for j in range(n_dest)]
        df_audit["Nearest_Destination"] = [dest_names[idx] for idx in closest_dest_idx]
        df_audit["Nearest_Dest_Lat"] = [dest_lats[idx] for idx in closest_dest_idx]
        df_audit["Nearest_Dest_Lon"] = [dest_lons[idx] for idx in closest_dest_idx]
        df_audit["Estimated_Travel_Distance_km"] = simulated_distances
        
        # Categorize accessibility
        # Green / High: <= 3 km. Orange / Moderate: 3 to 7 km. Red / Isolated: > 7 km.
        conditions = [
            df_audit["Estimated_Travel_Distance_km"] <= 3.0,
            (df_audit["Estimated_Travel_Distance_km"] > 3.0) & (df_audit["Estimated_Travel_Distance_km"] <= 7.0),
            df_audit["Estimated_Travel_Distance_km"] > 7.0
        ]
        choices = ["High Accessibility", "Moderate Accessibility", "Isolated (Transit Desert)"]
        df_audit["Accessibility_Status"] = np.select(conditions, choices, default="Moderate")
        
        # Correlate demographics with access
        # Find all numerical columns excluding lat, lon, calculations
        exclude = [start_lat, start_lon, "Nearest_Destination_Index", "Nearest_Dest_Lat", "Nearest_Dest_Lon", "Estimated_Travel_Distance_km"]
        num_cols = [c for c in df_start.columns if pd.api.types.is_numeric_dtype(df_start[c]) and c not in exclude]
        
        correlations = []
        for col in num_cols:
            # Split demographic average for Highly Accessible vs Isolated neighborhoods
            avg_accessible = df_audit[df_audit["Accessibility_Status"] == "High Accessibility"][col].mean()
            avg_isolated = df_audit[df_audit["Accessibility_Status"] == "Isolated (Transit Desert)"][col].mean()
            
            avg_accessible = 0.0 if np.isnan(avg_accessible) else avg_accessible
            avg_isolated = 0.0 if np.isnan(avg_isolated) else avg_isolated
            
            # Simple inequality quotient: Isolated Avg / Accessible Avg
            inequality_ratio = avg_isolated / avg_accessible if avg_accessible != 0 else 0.0
            
            correlations.append({
                "Socio-Demographic Factor": col,
                "Average in High-Access Areas": avg_accessible,
                "Average in Isolated Areas": avg_isolated,
                "Inequality Ratio (Isolated/Accessible)": inequality_ratio
            })
            
        df_compare = pd.DataFrame(correlations)
        return df_compare, df_audit
    except Exception as e:
        print(f"Network audit failed: {e}")
        return None, f"Network mapping failed: {e}"

def generate_network_map(df_audit, df_dest, start_lat, start_lon, dest_lat, dest_lon, dest_label):
    """Draws a beautiful Folium map showing connecting travel paths color-coded by accessibility."""
    mean_lat = df_dest[dest_lat].mean()
    mean_lon = df_dest[dest_lon].mean()
    
    m = folium.Map(location=[mean_lat, mean_lon], zoom_start=12, tiles="CartoDB dark_matter")
    
    # Color map for routes
    status_colors = {
        "High Accessibility": "#10b981",       # Green
        "Moderate Accessibility": "#f97316",   # Orange
        "Isolated (Transit Desert)": "#ef4444" # Red
    }
    
    # 1. Plot starting tract points and routes to closest destinations
    for i, row in df_audit.iterrows():
        s_lat = row[start_lat]
        s_lon = row[start_lon]
        d_lat = row["Nearest_Dest_Lat"]
        d_lon = row["Nearest_Dest_Lon"]
        d_name = row["Nearest_Destination"]
        dist = row["Estimated_Travel_Distance_km"]
        status = row["Accessibility_Status"]
        
        color = status_colors.get(status, "#6b7280")
        
        # Add travel connection path (direct line representing route)
        folium.PolyLine(
            locations=[[s_lat, s_lon], [d_lat, d_lon]],
            color=color,
            weight=2,
            opacity=0.6,
            dash_array="5, 5" if status == "Isolated (Transit Desert)" else None
        ).add_to(m)
        
        # Add neighborhood centroid circle
        folium.CircleMarker(
            location=[s_lat, s_lon],
            radius=6,
            color=color,
            fill=True,
            fill_color=color,
            fill_opacity=0.8,
            weight=1.5,
            popup=f"Nearest Hub: {d_name}<br>Simulated Distance: {dist:.2f} km<br>Status: {status}"
        ).add_to(m)
        
    # 2. Plot vital Destination Hub markers (high contrast white/gold icon)
    for i, row in df_dest.iterrows():
        lat = row[dest_lat]
        lon = row[dest_lon]
        name = row[dest_label] if dest_label in df_dest.columns else f"Hub {i}"
        
        folium.Marker(
            location=[lat, lon],
            popup=f"<b>Destination Hub: {name}</b>",
            icon=folium.Icon(color="orange", icon="home")
        ).add_to(m)
        
    return m

def full_network_pipeline(file_start, file_dest, start_lat, start_lon, dest_lat, dest_lon, dest_label, circuity):
    """Executes loading, network distance mapping, comparative audits, and downloads."""
    if file_start is None or file_dest is None:
        return None, "Please upload both the Demographic Starting Points CSV and Destination Hubs CSV files.", pd.DataFrame(), None
        
    try:
        df_start = pd.read_csv(file_start.name)
        df_dest = pd.read_csv(file_dest.name)
        
        # Column checks
        for c in [start_lat, start_lon]:
            if c not in df_start.columns:
                return None, f"ERROR: Demographic column '{c}' not found! Check columns.", pd.DataFrame(), None
        for c in [dest_lat, dest_lon]:
            if c not in df_dest.columns:
                return None, f"ERROR: Destination column '{c}' not found! Check columns.", pd.DataFrame(), None
                
        df_start_clean = df_start.dropna(subset=[start_lat, start_lon]).copy()
        df_dest_clean = df_dest.dropna(subset=[dest_lat, dest_lon]).copy()
        
        df_compare, df_audit = run_network_equity_audit(
            df_start_clean, df_dest_clean, start_lat, start_lon, dest_lat, dest_lon, dest_label, circuity
        )
        
        if df_compare is None:
            # df_compare holds the error string
            return None, df_audit, pd.DataFrame(), None
            
        # Draw map
        map_obj = generate_network_map(
            df_audit, df_dest_clean, start_lat, start_lon, dest_lat, dest_lon, dest_label
        )
        
        # Save HTML map
        temp_map = tempfile.NamedTemporaryFile(delete=False, suffix=".html")
        map_obj.save(temp_map.name)
        
        isolated_count = len(df_audit[df_audit["Accessibility_Status"] == "Isolated (Transit Desert)"])
        total_count = len(df_audit)
        
        status_md = f"""
        ### 📊 Network Transit Equity Metrics:
        * **Total Starting Neighborhoods**: `{total_count}`
        * **Isolated Neighborhoods (Transit Deserts)**: `{isolated_count} ({isolated_count/total_count:.1%})`
        * **Highly Accessible Areas (Short Travel)**: `{len(df_audit[df_audit["Accessibility_Status"] == "High Accessibility"])}`
        * **Applied Road Winding Circuity Factor**: `{circuity:.2f}x`
        
        *Interpretation*: If the **Inequality Ratio** on the right is greater than 1.0, it indicates that isolated neighborhoods have a higher concentration of that demographic attribute than accessible zones, proving spatial polarization!
        """
        
        # CSV download path
        temp_csv = tempfile.NamedTemporaryFile(delete=False, suffix="_transit_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"Transit audit processing failed: {e}", pd.DataFrame(), None

# Custom styling (Monochrome / Indigo theme)
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, #6366f1 0%, #4f46e5 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(
        """
        # 🕸️ Route Equity & Spatial Network Analyzer
        ### Analyze spatial access by calculating simulated road-network routing (using circuity multipliers) from neighborhoods to vital resources. Identify geographic exclusion and audit transit deserts.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=4):
            with gr.Card():
                gr.Markdown("### 1. Upload Starting Tract Coordinates")
                file_start_input = gr.File(label="Upload Neighborhood Centroids CSV", file_types=[".csv"])
                with gr.Row():
                    start_lat_name = gr.Textbox(label="Neighborhood Lat Column", value="Latitude")
                    start_lon_name = gr.Textbox(label="Neighborhood Lon Column", value="Longitude")
                    
            with gr.Card():
                gr.Markdown("### 2. Upload Destination Hubs (Hospitals/Services)")
                file_dest_input = gr.File(label="Upload Vital Hub POIs CSV", file_types=[".csv"])
                with gr.Row():
                    dest_lat_name = gr.Textbox(label="Vital Hub Lat Column", value="Latitude")
                    dest_lon_name = gr.Textbox(label="Vital Hub Lon Column", value="Longitude")
                dest_lbl_name = gr.Textbox(label="Hub Name/Label Column", value="Name")
                
            with gr.Card():
                gr.Markdown("### 3. Route Settings")
                circuity_slider = gr.Slider(
                    minimum=1.0, maximum=2.0, value=1.3, step=0.05,
                    label="Urban Circuity Winding Multiplier (Simulate roads vs. crow flies)"
                )
                analyze_btn = gr.Button("Analyze Route Network Equity", variant="primary", elem_classes="btn-primary")
                
        with gr.Column(scale=6):
            with gr.Tabs():
                with gr.TabItem("🗺️ Dynamic Transit Network Map"):
                    map_output = gr.HTML(label="Leaflet Route Map Grid", value="<div style='text-align: center; padding: 50px; color: gray;'>Map will load here...</div>")
                    summary_output = gr.Markdown("Please load coordinates and calculate routing.")
                    
                with gr.TabItem("📊 Socio-Spatial Inequality Report"):
                    table_output = gr.Dataframe(
                        label="Calculated Access Disparities Table (Isolated vs. High-Access)",
                        interactive=False,
                        wrap=True
                    )
                    download_btn = gr.File(label="Download Labeled CSV Database", interactive=False)

    analyze_btn.click(
        fn=full_network_pipeline,
        inputs=[file_start_input, file_dest_input, start_lat_name, start_lon_name, dest_lat_name, dest_lon_name, dest_lbl_name, circuity_slider],
        outputs=[map_output, summary_output, table_output, download_btn]
    )

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