#!/usr/bin/env python3
"""Resample DEM and population density to 100m grid."""

import pandas as pd
import rasterio
from rasterio.warp import transform


def resample_dem(dem_path: str, grid_parquet: str, output_path: str):
    import numpy as np
    df = pd.read_parquet(grid_parquet)

    with rasterio.open(dem_path) as src:
        elevations = []
        for lon, lat in zip(df['center_lon'], df['center_lat']):
            py, px = src.index(lon, lat)
            if 0 <= py < src.height and 0 <= px < src.width:
                elevations.append(src.read(1)[py, px])
            else:
                elevations.append(np.nan)

    result = pd.DataFrame({'grid_id': df['grid_id'], 'elevation_m': elevations})
    result.to_parquet(output_path, index=False)
    print(f"DEM saved: {output_path}")


def resample_population(pop_dir: str, grid_parquet: str, output_path: str):
    import numpy as np
    import glob
    df = pd.read_parquet(grid_parquet)

    tif_files = glob.glob(f"{pop_dir}/*.tif")
    if not tif_files:
        print(f"No TIF files found in {pop_dir}")
        return

    populations = []
    for lon, lat in zip(df['center_lon'], df['center_lat']):
        val = 0
        for tif_file in tif_files:
            try:
                with rasterio.open(tif_file) as src:
                    py, px = src.index(lon, lat)
                    if 0 <= py < src.height and 0 <= px < src.width:
                        val += src.read(1)[py, px]
            except:
                pass
        populations.append(val)

    result = pd.DataFrame({'grid_id': df['grid_id'], 'population_density': populations})
    result.to_parquet(output_path, index=False)
    print(f"Population saved: {output_path}")


def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--dem', default='Datas/DEM/CJJJD_DEM.TIF')
    parser.add_argument('--pop-dir', default='Datas/landscan-hd-china-v1-assets')
    parser.add_argument('--grid-parquet', default='processed/grid_100m_index.parquet')
    parser.add_argument('--output-dem', default='processed/grid_dem.parquet')
    parser.add_argument('--output-pop', default='processed/grid_population.parquet')
    args = parser.parse_args()

    print("Resampling DEM...")
    resample_dem(args.dem, args.grid_parquet, args.output_dem)

    print("Resampling population density...")
    resample_population(args.pop_dir, args.grid_parquet, args.output_pop)


if __name__ == '__main__':
    main()