Population from GHS

admin_id = "ITF2"  # Example admin ID for Central Greece

# Set working directory
workdir = "/etc/ecmwf/nfs/dh2_home_a/nejk/code/extreme_precip_exposure"

# Choose dataset: 'GHS_POP' (1975-2030) or 'GHS_WUP_POP' (1975-2100)
dataset_choice = "GHS_POP"  # Change to "GHS_WUP_POP" for urban population projections

import xarray as xr
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import geopandas as gpd
import regionmask
from pathlib import Path
import rioxarray as rxr
import os
import re

# Set working directoy
os.chdir(workdir) 
print(f"Working directory: {os.getcwd()}")

# Set up data directories
data_dir = Path("./data")
ghs_pop_dir = data_dir / "population" / "GHS_POP"
ghs_wup_dir = data_dir / "population" / "GHS_WUP_POP"

# Set directory and file pattern based on dataset choice
if dataset_choice == "GHS_POP":
    ghs_dir = ghs_pop_dir
    file_pattern = "*/GHS_POP_E*_GLOBE_R2023A_4326_30ss_V1_0.tif"
    dataset_name = "GHS-POP"
    year_range = "(1975-2030)"
elif dataset_choice == "GHS_WUP_POP":
    ghs_dir = ghs_wup_dir
    file_pattern = "*/GHS_WUP_POP_E*_GLOBE_R2025A_4326_30ss_V1_0.tif"
    dataset_name = "GHS-WUP"
    year_range = "(1975-2100)"
else:
    raise ValueError(f"Invalid dataset_choice: {dataset_choice}. Must be 'GHS_POP' or 'GHS_WUP_POP'")

print(f"\nSelected dataset: {dataset_name} {year_range}")
print(f"Data directory: {ghs_dir}")

# Read NUTS shapefiles
regions_dir = data_dir / 'regions'
nuts_shp = regions_dir / 'NUTS_RG_20M_2024_4326' / 'NUTS_RG_20M_2024_4326.shp'
nuts_gdf = gpd.read_file(nuts_shp)

# Select the region of interest
sel_gdf = nuts_gdf[nuts_gdf['NUTS_ID'] == admin_id]
print(f"Found {admin_id} region: {sel_gdf['NUTS_NAME'].values[0]}")
print(f"Bounding box: {sel_gdf.geometry.total_bounds}")
lon_min, lat_min, lon_max, lat_max = sel_gdf.geometry.total_bounds

# Create a regionmask from the admin region geometry
admin_mask = regionmask.from_geopandas(sel_gdf, names='NUTS_ID')

# List all GHS tif files based on selected dataset
ghs_files = sorted(list(ghs_dir.glob(file_pattern)))
print(f"Found {len(ghs_files)} {dataset_name} files:")
for f in ghs_files:
    print(f"  {f.name}")

# Read all GHS files and extract years (clipped to region bounding box)
datasets = []
years = []

for tif_file in ghs_files:
    # Read using rioxarray
    ds = rxr.open_rasterio(tif_file)
    
    # Clip using sel() instead of rio.clip_box to avoid artifacts
    # Note: y coordinates are typically in descending order in raster data
    ds_clipped = ds.sel(
        x=slice(lon_min, lon_max),
        y=slice(lat_max, lat_min)  # Note: reversed for descending y-coords
    )
    
    # Extract year from filename (looking for pattern E{YYYY})
    year_match = re.search(r'E(\d{4})', tif_file.name)
    if year_match:
        year = int(year_match.group(1))
        years.append(year)
        print(f"Processing year: {year}")
    else:
        print(f"Warning: Could not extract year from {tif_file.name}")
        continue
    
    datasets.append(ds_clipped)

# Sort datasets and years together by year
sorted_pairs = sorted(zip(years, datasets), key=lambda x: x[0])
years, datasets = zip(*sorted_pairs)

# Concatenate along new time dimension
pop_ds = xr.concat(datasets, dim="time")
pop_ds = pop_ds.rename("population")
pop_ds = pop_ds.squeeze("band", drop=True)

# Assign time based on extracted years
pop_ds["time"] = pd.to_datetime([f"{year}-01-01" for year in years])
pop_ds = pop_ds.assign_coords(time=pop_ds.time)

# Replace potential fill values with NaN (common fill values: -200, -99999, etc.)
pop_ds = pop_ds.where(pop_ds >= 0)

print(f"\n{dataset_name} dataset shape: {pop_ds.shape}")
print(f"Years available: {years}")
print(f"Coordinate ranges: x=[{float(pop_ds.x.min()):.2f}, {float(pop_ds.x.max()):.2f}], y=[{float(pop_ds.y.min()):.2f}, {float(pop_ds.y.max()):.2f}]")
print(f"Data value range: [{float(pop_ds.min()):.2f}, {float(pop_ds.max()):.2f}]")

# Data is already clipped to bounding box, so we use it directly as pop_region
pop_region = pop_ds

# Create admin region mask
pop_admin_mask_raw = admin_mask.mask(pop_region.x, pop_region.y)
pop_admin_mask = xr.DataArray(
    (~np.isnan(pop_admin_mask_raw.values)).astype(float),
    coords={'y': pop_region.y, 'x': pop_region.x},
    dims=['y', 'x']
)

print(f"{dataset_name} admin mask: {pop_admin_mask.sum().values:.0f} grid cells")

# Diagnostic: Check data quality 
print("Data diagnostics:")
print(f"  Population grid shape: {pop_region.shape}")
print(f"  Mask shape: {pop_admin_mask.shape}")
print(f"  Number of masked cells: {pop_admin_mask.sum().values:.0f}")
print(f"  Population data - NaN count: {pop_region.isnull().sum().values}")
print(f"  Population data - Valid values: {(~pop_region.isnull()).sum(['x', 'y']).values}")

# Check a single time slice
sample_time = 0
sample_data = pop_region.isel(time=sample_time)
print(f"\nSample year {pop_region.time.dt.year.values[sample_time]}:")
print(f"  Valid cells (not NaN): {(~sample_data.isnull()).sum().values}")
print(f"  Min value: {float(sample_data.min()):.2f}")
print(f"  Max value: {float(sample_data.max()):.2f}")
print(f"  Mean value: {float(sample_data.mean()):.2f}")

# Calculate total population for each year
pop_results = []

for i, year in enumerate(years):
    # Get data for this year
    data = pop_region.isel(time=i)
    
    # Apply mask and sum population
    masked_data = data.where(pop_admin_mask == 1, 0)
    total_pop = masked_data.sum().values
    
    pop_results.append({
        'year': year,
        'population': total_pop
    })
    
    print(f"{year}: {total_pop:,.0f} people")

df_pop = pd.DataFrame(pop_results)
print(f"\n{dataset_name} results:")
print(df_pop)

# Create output directory
output_dir = data_dir / admin_id / "ghs_population"
output_dir.mkdir(parents=True, exist_ok=True)

# Save results with dataset-specific filename
dataset_suffix = "ghs_pop" if dataset_choice == "GHS_POP" else "ghs_wup"
csv_file = output_dir / f"population_{dataset_suffix}_{admin_id}.csv"
df_pop.to_csv(csv_file, index=False)
print(f"Saved {dataset_name} results to: {csv_file}")