import os
import geopandas as gpd
import pandas as pd
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
import folium
from shapely.geometry import Point, LineString
from scipy.spatial import cKDTree
import warnings

warnings.filterwarnings('ignore')

# =========================================================
# SETTINGS
# =========================================================
FILES = {
    "routes": "busroute.geojson.geojsonl",
    "stops": "stops.geojson.geojsonl",
    "population": "population.geojson.geojsonl",
    "roads": "roads.geojson.geojsonl"
}

TARGET_CRS = "EPSG:3857"
WALK_BUFFER = 400 
ITERATIONS = 20

class FortWorthOptimizer:
    def __init__(self):
        print("--- Loading Fort Worth Data ---")
        self.routes = gpd.read_file(FILES["routes"], driver="GeoJSONSeq").to_crs(TARGET_CRS)
        self.stops = gpd.read_file(FILES["stops"], driver="GeoJSONSeq").to_crs(TARGET_CRS)
        self.pop_gdf = gpd.read_file(FILES["population"], driver="GeoJSONSeq").to_crs(TARGET_CRS)
        
        # Determine the "Study Area Bounds" based on population polygons
        # This will be used to crop the map and ignore outliers
        self.bounds = self.pop_gdf.total_bounds # [minx, miny, maxx, maxy]

        # Fix IDs
        if 'route_id' not in self.routes.columns:
            self.routes['route_id'] = self.routes.index.astype(str)
        self.routes['route_id'] = self.routes['route_id'].fillna("Unknown").astype(str)

        # Population numeric fix
        pop_col = next((c for c in ['population','pop','DN','VALUE'] if c in self.pop_gdf.columns), None)
        if pop_col:
            self.pop_gdf['pop_val'] = pd.to_numeric(self.pop_gdf[pop_col], errors='coerce').fillna(0)
        else:
            self.pop_gdf['pop_val'] = self.pop_gdf.geometry.area / 1000
        
        self.total_pop = self.pop_gdf['pop_val'].sum()

        print(f"--- Building Road Graph ({FILES['roads']}) ---")
        roads_gdf = gpd.read_file(FILES["roads"], driver="GeoJSONSeq").to_crs(TARGET_CRS)
        
        # DATA CLEANING: Only use roads within or near the study area to prevent "The Long Line"
        # We buffer the study area by 5km to allow for external connections
        study_area_buffered = self.pop_gdf.geometry.union_all().envelope.buffer(5000)
        roads_gdf = roads_gdf[roads_gdf.geometry.intersects(study_area_buffered)]

        self.road_graph = nx.Graph()
        for geom in roads_gdf.geometry:
            if isinstance(geom, LineString):
                coords = list(geom.coords)
                for i in range(len(coords) - 1):
                    u, v = coords[i], coords[i+1]
                    self.road_graph.add_edge(u, v, weight=Point(u).distance(Point(v)))
        
        self.nodes_list = list(self.road_graph.nodes)
        self.tree = cKDTree(self.nodes_list)
        print(f"Graph Ready: {len(self.nodes_list)} intersections within study area.")

    def snap_to_road(self, point):
        dist, idx = self.tree.query([point.x, point.y])
        return self.nodes_list[idx]

    def propose_route(self, current_stops, step):
        buffer = current_stops.geometry.buffer(WALK_BUFFER).union_all()
        underserved = self.pop_gdf[~self.pop_gdf.geometry.intersects(buffer)].copy()
        
        if underserved.empty: return None, None

        target_pt = underserved.sort_values(by='pop_val', ascending=False).iloc[0].geometry.centroid
        start_node = self.snap_to_road(target_pt)
        
        stop_coords = [(p.x, p.y) for p in current_stops.geometry]
        stop_tree = cKDTree(stop_coords)
        _, stop_idx = stop_tree.query([target_pt.x, target_pt.y])
        nearest_stop_pt = current_stops.iloc[stop_idx].geometry
        end_node = self.snap_to_road(nearest_stop_pt)

        try:
            path_coords = nx.shortest_path(self.road_graph, start_node, end_node, weight='weight')
            route_line = LineString(path_coords)
        except:
            return None, None

        rid = f"ROAD_ROUTE_{step}"
        new_route_gdf = gpd.GeoDataFrame([{'route_id': rid, 'geometry': route_line}], crs=TARGET_CRS)
        
        new_stops = []
        for d in np.arange(0, route_line.length, 400):
            new_stops.append({'route_id': rid, 'geometry': route_line.interpolate(d)})
        
        return new_route_gdf, gpd.GeoDataFrame(new_stops, crs=TARGET_CRS)

    def run(self):
        curr_routes, curr_stops = self.routes.copy(), self.stops.copy()
        results = []

        for i in range(ITERATIONS + 1):
            # Calculate Coverage
            buffer = curr_stops.geometry.buffer(WALK_BUFFER).union_all()
            res_temp = self.pop_gdf.copy()
            res_temp['orig_area'] = res_temp.geometry.area
            inter = gpd.overlay(res_temp, gpd.GeoDataFrame(geometry=[buffer], crs=TARGET_CRS), how='intersection')
            inter['served'] = (inter.geometry.area / inter['orig_area']) * inter['pop_val']
            cov = (inter['served'].sum() / self.total_pop) * 100
            
            print(f"Step {i}: Coverage {cov:.2f}%")
            results.append({'step': i, 'coverage': cov})

            # Plotting with AUTO-ZOOM
            fig, ax = plt.subplots(figsize=(10, 10))
            self.pop_gdf.plot(ax=ax, color='#eeeeee', edgecolor='#bcbcbc', label='Population')
            curr_routes.plot(ax=ax, color='blue', linewidth=0.8, alpha=0.3)
            
            # Draw proposed in red
            is_new = curr_routes['route_id'].str.contains("ROAD", na=False)
            if is_new.any():
                curr_routes[is_new].plot(ax=ax, color='red', linewidth=2.5)
            
            # --- THE FIX: CROP TO POPULATION BOUNDS ---
            ax.set_xlim([self.bounds[0], self.bounds[2]])
            ax.set_ylim([self.bounds[1], self.bounds[3]])
            
            plt.title(f"Fort Worth Optimization {i} | Coverage: {cov:.2f}%")
            plt.savefig(f"step_{i}.png", dpi=200, bbox_inches='tight')
            plt.close()

            if i < ITERATIONS:
                new_r, new_s = self.propose_route(curr_stops, i)
                if new_r is not None:
                    curr_routes = pd.concat([curr_routes, new_r], ignore_index=True)
                    curr_stops = pd.concat([curr_stops, new_s], ignore_index=True)

        curr_routes.to_file("proposed_road_network.geojson", driver="GeoJSON")
        print("Success. Maps are now zoomed to the city area.")

if __name__ == "__main__":
    FortWorthOptimizer().run()
