#!/usr/bin/env python3
"""Generate 100x100m resolution grid index for Wuhan city, China."""

import geopandas as gpd
import pandas as pd
import numpy as np
from shapely.geometry import box
from pathlib import Path


def generate_wuhan_grid(
    boundary_path: str = "Datas/武汉市.geojson",
    output_dir: str = "processed",
    grid_size: float = 100,
) -> tuple[gpd.GeoDataFrame, pd.DataFrame]:
    """Generate 100m resolution grid covering Wuhan boundary."""
    print(f"Loading Wuhan boundary from {boundary_path}...")
    wuhan = gpd.read_file(boundary_path)

    bounds = wuhan.total_bounds
    print(f"Wuhan bounds: minx={bounds[0]:.4f}, miny={bounds[1]:.4f}, maxx={bounds[2]:.4f}, maxy={bounds[3]:.4f}")

    minx, miny, maxx, maxy = bounds
    cell_size_deg = grid_size / 111000.0

    print(f"Creating grid with {grid_size}m cells (vectorized)...")

    x_coords = np.arange(minx, maxx, cell_size_deg)
    y_coords = np.arange(miny, maxy, cell_size_deg)
    print(f"  Grid dimensions: {len(x_coords)} x {len(y_coords)}")

    x_grid, y_grid = np.meshgrid(x_coords, y_coords)
    x_flat = x_grid.flatten()
    y_flat = y_grid.flatten()

    print(f"  Total cells in bounding box: {len(x_flat)}")

    minxs = x_flat
    minys = y_flat
    maxxs = minxs + cell_size_deg
    maxys = minys + cell_size_deg

    geometries = [box(mx, my, Mx, My) for mx, my, Mx, My in zip(minxs, minys, maxxs, maxys)]

    cells = np.arange(len(geometries))
    rows = cells // len(x_coords)
    cols = cells % len(x_coords)

    print("  Building GeoDataFrame...")
    grid_gdf = gpd.GeoDataFrame({
        'row': rows,
        'col': cols,
        'geometry': geometries
    }, crs="EPSG:4326")

    print("Filtering to cells intersecting Wuhan boundary...")
    wuhan_union = wuhan.unary_union
    mask = grid_gdf.intersects(wuhan_union)
    grid_gdf = grid_gdf[mask].copy().reset_index(drop=True)

    print(f"Cells within Wuhan boundary: {len(grid_gdf)}")

    grid_gdf['grid_id'] = [f"r{r}_c{c}" for r, c in zip(grid_gdf['row'], grid_gdf['col'])]

    centroids = grid_gdf.geometry.centroid
    grid_gdf['center_lon'] = centroids.x
    grid_gdf['center_lat'] = centroids.y
    grid_gdf['polygon'] = grid_gdf.geometry.apply(lambda g: g.wkt)

    parquet_df = grid_gdf[['grid_id', 'center_lon', 'center_lat', 'row', 'col', 'polygon']].copy()

    return grid_gdf, parquet_df


def main():
    output_dir = Path("processed")
    output_dir.mkdir(parents=True, exist_ok=True)

    grid_gdf, parquet_df = generate_wuhan_grid()

    geojson_path = output_dir / "grid_100m_index.geojson"
    print(f"Exporting to GeoJSON: {geojson_path}")
    grid_gdf.to_file(geojson_path, driver="GeoJSON")
    print(f"  Exported {len(grid_gdf)} features")

    parquet_path = output_dir / "grid_100m_index.parquet"
    print(f"Exporting to Parquet: {parquet_path}")
    parquet_df.to_parquet(parquet_path, index=False)
    print(f"  Exported {len(parquet_df)} rows")

    print("\n=== Grid Summary ===")
    print(f"Total grid cells: {len(grid_gdf)}")
    print(f"Bounds: {grid_gdf.total_bounds}")
    print(f"Grid ID format example: {grid_gdf['grid_id'].iloc[0]}")
    print(f"Center coordinate range:")
    print(f"  Lon: {parquet_df['center_lon'].min():.4f} to {parquet_df['center_lon'].max():.4f}")
    print(f"  Lat: {parquet_df['center_lat'].min():.4f} to {parquet_df['center_lat'].max():.4f}")


if __name__ == "__main__":
    main()