#!/usr/bin/env python3
"""
PostGIS index optimization script for grid queries.

Creates spatial indexes on the grids table for efficient bounding box
and radius queries used by the monitoring and prediction APIs.

Usage:
    python scripts/setup_postgis_indexes.py --connection postgresql://user:pass@localhost:5432/wuhan_disease
"""

import sys
import argparse

def create_indexes(connection_string):
    import asyncpg

    indexes = [
        ("idx_grids_geometry", "CREATE INDEX IF NOT EXISTS idx_grids_geometry ON grids USING GIST (geometry)"),
        ("idx_grids_centroid", "CREATE INDEX IF NOT EXISTS idx_grids_centroid ON grids USING GIST (ST_Transform(geometry, 32650))"),
        ("idx_grids_grid_id", "CREATE INDEX IF NOT EXISTS idx_grids_grid_id ON grids (grid_id)"),
        ("idx_grids_district", "CREATE INDEX IF NOT EXISTS idx_grids_district ON grids (district)"),
    ]

    print("Creating PostGIS spatial indexes...")

    async def run_indexes():
        conn = await asyncpg.connect(connection_string)

        for idx_name, sql in indexes:
            try:
                await conn.execute(sql)
                print(f"  Created: {idx_name}")
            except Exception as e:
                print(f"  Failed: {idx_name} - {e}")

        await conn.close()

    import asyncio
    asyncio.run(run_indexes())
    print("Done!")


def create_grid_table_sql():

    return """
-- Create grids table for 100m grid cells
CREATE TABLE IF NOT EXISTS grids (
    grid_id VARCHAR(20) PRIMARY KEY,
    geometry GEOMETRY(POLYGON, 4326) NOT NULL,
    center_lat DOUBLE PRECISION NOT NULL,
    center_lon DOUBLE PRECISION NOT NULL,
    district VARCHAR(50),
    dem DOUBLE PRECISION,
    population_density DOUBLE PRECISION,
    created_at TIMESTAMP DEFAULT NOW()
);

-- Spatial index for geometry queries
CREATE INDEX IF NOT EXISTS idx_grids_geometry ON grids USING GIST (geometry);

-- Index for district lookups
CREATE INDEX IF NOT EXISTS idx_grids_district ON grids (district);

-- Index for grid_id lookups
CREATE INDEX IF NOT EXISTS idx_grids_grid_id ON grids (grid_id);

-- Index for bounding box queries (UTM projection for meters)
CREATE INDEX IF NOT EXISTS idx_grids_centroid ON grids USING GIST (ST_Transform(geometry, 32650));

-- Cluster the table by geometry for better spatial query performance
CLUSTER grids USING idx_grids_geometry;

-- Analyze the table for query planner
ANALYZE grids;

-- Example queries:

-- 1. Bounding box query (within 114.0-115.0 lon, 29.5-30.5 lat)
SELECT grid_id, center_lat, center_lon
FROM grids
WHERE geometry && ST_MakeEnvelope(113.8, 29.4, 115.2, 30.6, 4326);

-- 2. Radius query (within 10km of point)
SELECT grid_id, center_lat, center_lon,
       ST_Distance(geometry, ST_Transform(ST_SetSRID(ST_MakePoint(114.3, 30.6), 4326), 32650)) as distance
FROM grids
WHERE ST_DWithin(
    ST_Transform(geometry, 32650),
    ST_Transform(ST_SetSRID(ST_MakePoint(114.3, 30.6), 4326), 32650),
    10000
)
ORDER BY distance
LIMIT 100;

-- 3. District aggregation
SELECT district, COUNT(*) as grid_count, AVG(population_density) as avg_pop
FROM grids
GROUP BY district
ORDER BY grid_count DESC;
"""


def main():
    parser = argparse.ArgumentParser(description='PostGIS index setup for grid queries')
    parser.add_argument('--connection', help='PostgreSQL connection string')
    parser.add_argument('--sql-only', action='store_true', help='Print SQL only')
    args = parser.parse_args()

    if args.sql_only:
        print(create_grid_table_sql())
    elif args.connection:
        create_indexes(args.connection)
    else:
        print("Usage:")
        print("  python scripts/setup_postgis_indexes.py --sql-only  # Print SQL")
        print("  python scripts/setup_postgis_indexes.py --connection postgresql://...")


if __name__ == '__main__':
    main()