# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_sidecar.ipynb.

# %% auto 0
__all__ = ['logger', 'PSF_REGISTRY', 'compute_healpix_ids_from_lonlat', 'detect_lonlat_columns', 'compute_geo_statistics',
           'format_geo_statistics', 'process_partition', 'add_psf_weights_to_sidecar', 'normalize_weights_per_cell',
           'compute_assignment_weight', 'build_output_path', 'write_sidecar_metadata', 'validate_nside',
           'parse_arguments', 'PSF', 'GaussianPSF', 'get_psf', 'write_partitioned_output', 'write_coalesced_output',
           'main', 'get_healpix_cell_geometry']

# %% ../nbs/01_sidecar.ipynb 2
#| eval: true
"""healpix_sidecar.py

Create a lightweight sidecar (index-aside) that maps source geometries to HEALPix cells.

Requirements:
- dask_geopandas for lazy/parallel geoparquet reading
- cdshealpix for HEALPix computation (falls back to healpy if unavailable)
- shapely for coordinate extraction
- argparse for CLI

Output: parquet file with only two columns: source_id (original index) and healpix_id (uint64)

Usage example:
  python healpix_sidecar.py --input data.parquet --nside 64 --mode fuzzy --ncores 8
"""
import argparse
import logging
import os
from pathlib import Path
import sys
import numpy as np

import pandas as pd

try:
    import dask_geopandas as dg
except Exception:
    raise ImportError("This script requires dask_geopandas. Install it with `pip install dask-geopandas`.")

try:
    # cdshealpix is preferred for speed
    import cdshealpix as ch
except Exception:
    ch = None

try:
    import healpy as _healpy
except Exception:
    _healpy = None

from shapely import get_coordinates  # shapely>=2.0
from shapely.geometry import Polygon, MultiPolygon
from tqdm.auto import tqdm
import antimeridian

try:
    from dask.diagnostics.progress import ProgressBar as DaskProgressBar
    DASK_PROGRESS_AVAILABLE = True
except ImportError:
    try:
        from dask.diagnostics import ProgressBar as DaskProgressBar
        DASK_PROGRESS_AVAILABLE = True
    except ImportError:
        DASK_PROGRESS_AVAILABLE = False
        DaskProgressBar = None

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("healpix_sidecar")

# %% ../nbs/01_sidecar.ipynb 4
def compute_healpix_ids_from_lonlat(nside: int, lons: np.ndarray, lats: np.ndarray) -> np.ndarray:
    """Compute HEALPix indices for arrays of lon,lat in degrees.

    Tries to use cdshealpix if available, otherwise falls back to healpy.
    Returns a 1D integer numpy array of same length as inputs.
    """
    if lons.size == 0:
        return np.array([], dtype=np.int64)

    # normalize lons to [0,360)
    lons = np.mod(lons.astype(float), 360.0)
    lats = lats.astype(float)

    # Prefer healpy (to match notebook usage). Fall back to cdshealpix if healpy not available.
    if _healpy is not None:
        # healpy expects theta (colat) and phi (lon) in radians
        phi = np.radians(lons)
        theta = np.radians(90.0 - lats)
        return _healpy.ang2pix(nside, theta, phi, nest=True)

    if ch is not None:
        try:
            return np.asarray(ch.lonlat_to_healpix(nside, lons, lats, nest=True), dtype=np.int64)
        except Exception:
            try:
                return np.asarray(ch.lonlat_to_healpix(nside, lons, lats), dtype=np.int64)
            except Exception:
                logger.debug("cdshealpix present but call failed")

    raise RuntimeError("No HEALPix implementation available: install healpy or cdshealpix")

# %% ../nbs/01_sidecar.ipynb 5
def detect_lonlat_columns(gdf_sample) -> tuple[str | None, str | None]:
    """Auto-detect longitude and latitude columns from a GeoDataFrame sample.
    
    Returns:
        Tuple of (lon_column, lat_column) or (None, None) if not found
    """
    # Common column name patterns (case-insensitive)
    lon_patterns = ['lon', 'longitude', 'long', 'x', 'easting']
    lat_patterns = ['lat', 'latitude', 'y', 'northing']
    
    cols_lower = {col.lower(): col for col in gdf_sample.columns if col != 'geometry'}
    
    lon_col = None
    lat_col = None
    
    # Try to find longitude column
    for pattern in lon_patterns:
        for col_lower, col_orig in cols_lower.items():
            if pattern in col_lower:
                lon_col = col_orig
                break
        if lon_col:
            break
    
    # Try to find latitude column
    for pattern in lat_patterns:
        for col_lower, col_orig in cols_lower.items():
            if pattern in col_lower:
                lat_col = col_orig
                break
        if lat_col:
            break
    
    return lon_col, lat_col


def compute_geo_statistics(input_path: Path, lon_col: str | None = None, 
                          lat_col: str | None = None, 
                          sample_size: int = 10000,
                          lon_convention: str | None = None) -> dict:
    """Compute geographical statistics for a GeoParquet file using DuckDB for efficiency.
    
    This function can analyze raw data or apply filtering based on longitude convention.
    
    Args:
        input_path: Path to input GeoParquet file
        lon_col: Name of longitude column (if None, will auto-detect or extract from geometry)
        lat_col: Name of latitude column (if None, will auto-detect or extract from geometry)
        sample_size: Number of rows to sample for geometry-based extraction (if needed)
        lon_convention: Optional longitude convention for filtering:
                       '0_360' for [0,360) × [-90,90]
                       'minus_plus180' for [-180,180) × [-90,90]
                       None (default) for no filtering (raw data)
        
    Returns:
        Dictionary with statistics: {
            'lon': {'min', 'max', 'mean', 'std', 'count'},
            'lat': {'min', 'max', 'mean', 'std', 'count'},
            'source': 'columns' or 'geometry',
            'lon_col': column name or None,
            'lat_col': column name or None,
            'filtered': bool (True if convention filtering was applied),
            'total_count': int (total records before filtering, if filtered),
            'filtered_count': int (records after filtering, if filtered)
        }
    """
    import duckdb
    import geopandas as gpd
    
    logger.info(f"Computing geo-statistics for {input_path}")
    
    # Read a small sample to detect columns
    # Try GeoParquet first, fallback to regular parquet
    try:
        sample_gdf = gpd.read_parquet(input_path, max_rows=100)
    except Exception as e:
        logger.debug(f"Could not read as GeoParquet: {e}")
        try:
            # Fallback to pandas for regular parquet using pyarrow
            import pyarrow.parquet as pq
            table = pq.read_table(input_path, columns=None)
            sample_df = table.slice(0, 100).to_pandas()
            # Create a mock GeoDataFrame structure for column detection
            sample_gdf = sample_df
        except Exception as e2:
            logger.warning(f"Could not read parquet file: {e2}")
            return {}
    
    # Auto-detect columns if not provided
    if lon_col is None or lat_col is None:
        detected_lon, detected_lat = detect_lonlat_columns(sample_gdf)
        lon_col = lon_col or detected_lon
        lat_col = lat_col or detected_lat
    
    # Check if columns exist in the data
    has_lon_col = lon_col is not None and lon_col in sample_gdf.columns
    has_lat_col = lat_col is not None and lat_col in sample_gdf.columns
    
    result = {
        'lon': {},
        'lat': {},
        'source': None,
        'lon_col': lon_col,
        'lat_col': lat_col,
        'filtered': lon_convention is not None,
        'lon_convention': lon_convention
    }
    
    # Determine filtering bounds based on convention
    where_clause = ""
    if lon_convention == '0_360':
        where_clause = f'WHERE "{lon_col}" >= 0 AND "{lon_col}" < 360 AND "{lat_col}" >= -90 AND "{lat_col}" <= 90'
        logger.info(f"Applying filtering: lon=[0,360), lat=[-90,90]")
    elif lon_convention == 'minus_plus180':
        where_clause = f'WHERE "{lon_col}" >= -180 AND "{lon_col}" < 180 AND "{lat_col}" >= -90 AND "{lat_col}" <= 90'
        logger.info(f"Applying filtering: lon=[-180,180), lat=[-90,90]")
    elif lon_convention is not None:
        logger.warning(f"Unknown lon_convention '{lon_convention}', computing raw statistics")
        result['filtered'] = False
    
    try:
        # Strategy 1: Use explicit lon/lat columns if available
        if has_lon_col and has_lat_col:
            logger.info(f"Using explicit columns: lon='{lon_col}', lat='{lat_col}'")
            result['source'] = 'columns'
            
            # Use DuckDB for efficient statistics computation
            con = duckdb.connect()
            
            # Get total count before filtering (if filtering is applied)
            if where_clause:
                total_query = f"SELECT COUNT(*) FROM read_parquet('{str(input_path)}')"
                total_count = con.execute(total_query).fetchone()[0]
                result['total_count'] = int(total_count)
            
            # Build query with optional WHERE clause for filtering
            query = f"""
            SELECT 
                COUNT(*) as count,
                MIN("{lon_col}") as lon_min,
                MAX("{lon_col}") as lon_max,
                AVG("{lon_col}") as lon_mean,
                STDDEV("{lon_col}") as lon_std,
                MIN("{lat_col}") as lat_min,
                MAX("{lat_col}") as lat_max,
                AVG("{lat_col}") as lat_mean,
                STDDEV("{lat_col}") as lat_std
            FROM read_parquet('{str(input_path)}')
            {where_clause}
            """
            
            stats = con.execute(query).fetchone()
            con.close()
            
            filtered_count = int(stats[0])
            
            result['lon'] = {
                'min': float(stats[1]) if stats[1] is not None else None,
                'max': float(stats[2]) if stats[2] is not None else None,
                'mean': float(stats[3]) if stats[3] is not None else None,
                'std': float(stats[4]) if stats[4] is not None else None,
                'count': filtered_count
            }
            result['lat'] = {
                'min': float(stats[5]) if stats[5] is not None else None,
                'max': float(stats[6]) if stats[6] is not None else None,
                'mean': float(stats[7]) if stats[7] is not None else None,
                'std': float(stats[8]) if stats[8] is not None else None,
                'count': filtered_count
            }
            
            if where_clause:
                result['filtered_count'] = filtered_count
            
        # Strategy 2: Extract from geometry column (sample-based for efficiency)
        else:
            logger.info("Extracting coordinates from geometry column (sampling)")
            result['source'] = 'geometry'
            
            # Read a larger sample for better statistics
            try:
                gdf = gpd.read_parquet(input_path, max_rows=sample_size)
            except Exception:
                gdf = sample_gdf
            
            # Extract coordinates from geometries
            coords_list = []
            for geom in gdf.geometry:
                if geom is None or geom.is_empty:
                    continue
                
                # Handle different geometry types
                if geom.geom_type == 'Point':
                    coords_list.append([geom.x, geom.y])
                elif geom.geom_type in ['Polygon', 'MultiPolygon']:
                    # Use centroid for polygons
                    centroid = geom.centroid
                    coords_list.append([centroid.x, centroid.y])
                else:
                    # Try to get any coordinates
                    try:
                        c = get_coordinates(geom)
                        if len(c) > 0:
                            coords_list.append([c[0, 0], c[0, 1]])
                    except Exception:
                        continue
            
            if coords_list:
                coords = np.array(coords_list)
                lons = coords[:, 0]
                lats = coords[:, 1]
                
                # Filter out invalid values
                valid_mask = np.isfinite(lons) & np.isfinite(lats)
                lons = lons[valid_mask]
                lats = lats[valid_mask]
                
                if len(lons) > 0:
                    result['lon'] = {
                        'min': float(np.min(lons)),
                        'max': float(np.max(lons)),
                        'mean': float(np.mean(lons)),
                        'std': float(np.std(lons)),
                        'count': len(lons)
                    }
                    result['lat'] = {
                        'min': float(np.min(lats)),
                        'max': float(np.max(lats)),
                        'mean': float(np.mean(lats)),
                        'std': float(np.std(lats)),
                        'count': len(lats)
                    }
                    
                    logger.info(f"Computed statistics from {len(lons)} sampled geometries")
    
    except Exception as e:
        logger.warning(f"Failed to compute geo-statistics: {e}")
        return {}
    
    return result


def format_geo_statistics(stats: dict) -> str:
    """Format geo-statistics for display using rich tables.
    
    Args:
        stats: Statistics dictionary from compute_geo_statistics
        
    Returns:
        Formatted string representation
    """
    if not stats or not stats.get('lon') or not stats.get('lat'):
        return "No geo-statistics available"
    
    try:
        from rich.console import Console
        from rich.table import Table
        from io import StringIO
        
        console = Console(file=StringIO(), width=100)
        
        # Create main statistics table
        table = Table(title="Geographical Statistics", show_header=True, header_style="bold magenta")
        table.add_column("Statistic", style="cyan", width=12)
        table.add_column("Longitude", justify="right", style="green")
        table.add_column("Latitude", justify="right", style="green")
        
        lon = stats['lon']
        lat = stats['lat']
        
        # Add rows
        table.add_row("Count", f"{lon.get('count', 'N/A'):,}", f"{lat.get('count', 'N/A'):,}")
        table.add_row("Min", f"{lon.get('min', float('nan')):.6f}", f"{lat.get('min', float('nan')):.6f}")
        table.add_row("Max", f"{lon.get('max', float('nan')):.6f}", f"{lat.get('max', float('nan')):.6f}")
        table.add_row("Mean", f"{lon.get('mean', float('nan')):.6f}", f"{lat.get('mean', float('nan')):.6f}")
        table.add_row("Std Dev", f"{lon.get('std', float('nan')):.6f}", f"{lat.get('std', float('nan')):.6f}")
        
        console.print(table)
        
        # Add metadata
        console.print(f"\n[bold]Data Source:[/bold] {stats.get('source', 'unknown')}")
        if stats.get('lon_col'):
            console.print(f"[bold]Longitude Column:[/bold] {stats['lon_col']}")
        if stats.get('lat_col'):
            console.print(f"[bold]Latitude Column:[/bold] {stats['lat_col']}")
        
        # Show filtering information
        if stats.get('filtered'):
            total = stats.get('total_count', 0)
            filtered = stats.get('filtered_count', 0)
            if total > 0:
                pct = 100.0 * filtered / total
                dropped = total - filtered
                console.print(f"\n[bold yellow]Filtering Applied:[/bold yellow] --lon-convention {stats.get('lon_convention')}")
                console.print(f"  Total records: {total:,}")
                console.print(f"  After filtering: {filtered:,} ({pct:.1f}%)")
                if dropped > 0:
                    console.print(f"  [red]Dropped: {dropped:,} ({100.0 - pct:.1f}%)[/red]")
        else:
            console.print(f"\n[bold]Filtering:[/bold] None (raw data)")
        
        # Validation warnings
        lon_min, lon_max = lon.get('min'), lon.get('max')
        lat_min, lat_max = lat.get('min'), lat.get('max')
        
        warnings = []
        
        if lat_min is not None and lat_max is not None:
            if lat_min < -90 or lat_max > 90:
                warnings.append(f"⚠️  Latitude out of valid range [-90, 90]: [{lat_min:.2f}, {lat_max:.2f}]")
        
        # Suggest appropriate longitude convention based on data
        if lon_min is not None and lon_max is not None:
            if lon_min >= 0 and lon_max <= 360:
                console.print(f"\n[bold cyan]Suggested convention:[/bold cyan] --lon-convention 0_360")
            elif lon_min >= -180 and lon_max <= 180:
                console.print(f"\n[bold cyan]Suggested convention:[/bold cyan] --lon-convention minus_plus180")
            else:
                warnings.append(f"⚠️  Longitude range [{lon_min:.2f}, {lon_max:.2f}] doesn't fit standard conventions")
        
        if warnings:
            console.print("\n[bold red]Validation Warnings:[/bold red]")
            for warning in warnings:
                console.print(f"  {warning}")
        
        # Get the string output
        output = console.file.getvalue()
        return output
        
    except ImportError:
        # Fallback to simple text formatting if rich is not available
        lines = ["=" * 60]
        lines.append("GEOGRAPHICAL STATISTICS")
        lines.append("=" * 60)
        lines.append(f"{'Statistic':<15} {'Longitude':>20} {'Latitude':>20}")
        lines.append("-" * 60)
        
        lon = stats['lon']
        lat = stats['lat']
        
        lines.append(f"{'Count':<15} {lon.get('count', 'N/A'):>20,} {lat.get('count', 'N/A'):>20,}")
        lines.append(f"{'Min':<15} {lon.get('min', float('nan')):>20.6f} {lat.get('min', float('nan')):>20.6f}")
        lines.append(f"{'Max':<15} {lon.get('max', float('nan')):>20.6f} {lat.get('max', float('nan')):>20.6f}")
        lines.append(f"{'Mean':<15} {lon.get('mean', float('nan')):>20.6f} {lat.get('mean', float('nan')):>20.6f}")
        lines.append(f"{'Std Dev':<15} {lon.get('std', float('nan')):>20.6f} {lat.get('std', float('nan')):>20.6f}")
        lines.append("=" * 60)
        lines.append(f"Data Source: {stats.get('source', 'unknown')}")
        if stats.get('lon_col'):
            lines.append(f"Longitude Column: {stats['lon_col']}")
        if stats.get('lat_col'):
            lines.append(f"Latitude Column: {stats['lat_col']}")
        
        # Show filtering information
        if stats.get('filtered'):
            total = stats.get('total_count', 0)
            filtered = stats.get('filtered_count', 0)
            if total > 0:
                pct = 100.0 * filtered / total
                dropped = total - filtered
                lines.append(f"\nFiltering Applied: --lon-convention {stats.get('lon_convention')}")
                lines.append(f"  Total records: {total:,}")
                lines.append(f"  After filtering: {filtered:,} ({pct:.1f}%)")
                if dropped > 0:
                    lines.append(f"  Dropped: {dropped:,} ({100.0 - pct:.1f}%)")
        else:
            lines.append("\nFiltering: None (raw data)")
        
        return "\n".join(lines)

# %% ../nbs/01_sidecar.ipynb 6
def process_partition(gdf, nside: int, mode: str, base_index: int | None = None, 
                     lon_convention: str = 'None',
                     data_psf=None, cell_psf=None, combine_method='multiply'
                        ) -> pd.DataFrame:
    """Process a single dask partition (GeoDataFrame) and return DataFrame of assignments.

    The returned D/ataFrame has columns ['source_id', 'healpix_id'] and one row per assignment
    (for strict mode: at most one row per source_id; for fuzzy mode: one row per touched healpix cell).
    
    Args:
        gdf: GeoDataFrame partition
        nside: HEALPix nside parameter
        mode: 'strict' or 'fuzzy' assignment mode
        base_index: Base index for source_id generation
        lon_convention: Longitude convention - '0_360' for [0,360) or 'minus_plus180' for [-180,180)
    """
    import pandas as _pd

    # Set lon/lat bounds based on convention
    if lon_convention == '0_360':
        lon_min, lon_max = 0.0, 360.0
        lat_min, lat_max = -90.0, 90.0
    elif lon_convention == 'minus_plus180':
        lon_min, lon_max = -180.0, 180.0
        lat_min, lat_max = -90.0, 90.0
    else:
        raise ValueError(f"Invalid lon_convention: {lon_convention}. Use '0_360' or 'minus_plus180'")

    logger.debug(f"Using lon_convention='{lon_convention}' with bounds: lon=[{lon_min}, {lon_max}), lat=[{lat_min}, {lat_max}]")

    out_rows = []
    dropped_count = 0
    total_count = 0
    dropped_prefilter = 0

    # if the partition is empty, return empty DataFrame
    if gdf is None or len(gdf) == 0:
        return _pd.DataFrame(columns=["source_id", "healpix_id"])

    # Determine source ids. If `base_index` is provided we generate sequential
    # global row numbers to match `gdf.reset_index()` semantics from the notebook.
    if base_index is not None:
        src_ids = base_index + np.arange(len(gdf), dtype=np.int64)
    else:
        # Prefer an explicit 'source_id' column if present; otherwise fall back to the index.
        if "source_id" in gdf.columns:
            src_ids = gdf["source_id"].to_numpy()
        else:
            src_ids = gdf.index.to_numpy()

    # iterate rows; keep work per-geometry contained to avoid large in-memory structures
    for src_id, geom in zip(src_ids, gdf.geometry.to_numpy()):
        try:
            # handle missing/empty geometries: match notebook filtering semantics by skipping
            # geometries that are None/empty rather than emitting NA rows
            if geom is None or geom.is_empty:
                continue

            # notebook pipeline: first validate the ORIGINAL geometry (pre-fix),
            # then apply antimeridian.fix_polygon and extract coordinates.
            def _is_valid_latitude(geometry):
                if geometry is None or geometry.is_empty:
                    return False
                geoms = [geometry] if getattr(geometry, "geom_type", "") == "Polygon" else list(getattr(geometry, "geoms", [geometry]))
                for g in geoms:
                    if getattr(g, "exterior", None) is not None:
                        for coord in g.exterior.coords:
                            lon = coord[0]
                            lat = coord[1]
                            if not (np.isfinite(lon) and np.isfinite(lat)):
                                return False
                            # check against configurable bounds
                            if lat < lat_min or lat > lat_max:
                                return False
                            if lon < lon_min or lon > lon_max:
                                return False
                    for interior in getattr(g, "interiors", []):
                        for coord in interior.coords:
                            lon = coord[0]
                            lat = coord[1]
                            if not (np.isfinite(lon) and np.isfinite(lat)):
                                return False
                            if lat < lat_min or lat > lat_max:
                                return False
                            if lon < lon_min or lon > lon_max:
                                return False
                return True

            total_count += 1
            # if original geometry fails the pre-fix filter, skip it entirely
            if not _is_valid_latitude(geom):
                dropped_count += 1
                dropped_prefilter += 1
                logger.debug(f"Dropped geometry {src_id} during pre-filter (lon_convention={lon_convention})")
                continue

            # apply the antimeridian fix (not used for the pre-filter decision)
            try:
                geom2 = antimeridian.fix_polygon(geom)
            except Exception as e:
                logger.debug(f"Antimeridian fix failed for {src_id}: {e}")
                geom2 = geom

            # vectored extraction using shapely.get_coordinates when available (on the fixed geometry)
            try:
                coords = get_coordinates(geom2)
            except Exception:
                # fallback gather exterior + interiors for polygons/multipolygons
                coords_list = []
                if isinstance(geom2, Polygon):
                    coords_list.extend(np.asarray(geom2.exterior.coords, dtype=float))
                    for r in geom2.interiors:
                        coords_list.extend(np.asarray(r.coords, dtype=float))
                elif isinstance(geom2, MultiPolygon):
                    for part in geom2.geoms:
                        coords_list.extend(np.asarray(part.exterior.coords, dtype=float))
                        for r in part.interiors:
                            coords_list.extend(np.asarray(r.coords, dtype=float))
                else:
                    # generic fallback: treat as single-vertex geometry
                    try:
                        coords_list = np.asarray(list(geom2.coords), dtype=float)
                    except Exception:
                        coords_list = []

                if len(coords_list) == 0:
                    # no coordinates to assign -> skip (matches notebook behaviour where empty lists
                    # later produce no exploded rows)
                    continue
                coords = np.asarray(coords_list, dtype=float)

            if coords.size == 0:
                continue

            # Extract lon/lat from coordinates (no normalization)
            lons = coords[:, 0].astype(float)
            lats = coords[:, 1].astype(float)

            # Filter invalid lat/lon values - filtering already done in pre-validation
            # This is just to catch any NaN/Inf values after antimeridian processing
            mask = np.isfinite(lons) & np.isfinite(lats)
            if not np.any(mask):
                dropped_count += 1
                logger.debug(f"Dropped geometry {src_id} - all coordinates non-finite after antimeridian fix")
                continue

            lons = lons[mask]
            lats = lats[mask]

            # compute healpix indices for all vertices
            hids = compute_healpix_ids_from_lonlat(nside, lons, lats)
            if hids.size == 0:
                continue

            unique = np.unique(hids)
            if mode == "strict":
                # only accept if all vertices fall in same single HEALPix cell
                if unique.size == 1:
                    weight = 1.0
                    if data_psf or cell_psf:
                        # For strict, use centroid of geometry and cell
                        src_centroid = geom.centroid
                        cell_geom = get_healpix_cell_geometry(unique[0], nside)
                        dx = src_centroid.x - cell_geom.centroid.x
                        dy = src_centroid.y - cell_geom.centroid.y
                        w_data = data_psf(dx, dy) if data_psf else 1.0
                        w_cell = cell_psf(dx, dy) if cell_psf else 1.0
                        if combine_method == 'multiply':
                            weight = w_data * w_cell
                        elif combine_method == 'sum':
                            weight = w_data + w_cell
                        elif combine_method == 'min':
                            weight = min(w_data, w_cell)
                        elif combine_method == 'max':
                            weight = max(w_data, w_cell)
                    out_rows.append({"source_id": int(src_id), "healpix_id": int(unique[0]), "weight": weight})
                else:
                    # If not all vertices in the same cell, still output with weight=1.0
                    out_rows.append({"source_id": int(src_id), "healpix_id": int(unique[0]), "weight": 1.0})
            else:
                # fuzzy: replicate source for each unique cell
                for hid in unique:
                    weight = 1.0
                    if data_psf or cell_psf:
                        src_centroid = geom.centroid
                        cell_geom = get_healpix_cell_geometry(hid, nside)
                        dx = src_centroid.x - cell_geom.centroid.x
                        dy = src_centroid.y - cell_geom.centroid.y
                        w_data = data_psf(dx, dy) if data_psf else 1.0
                        w_cell = cell_psf(dx, dy) if cell_psf else 1.0
                        if combine_method == 'multiply':
                            weight = w_data * w_cell
                        elif combine_method == 'sum':
                            weight = w_data + w_cell
                        elif combine_method == 'min':
                            weight = min(w_data, w_cell)
                        elif combine_method == 'max':
                            weight = max(w_data, w_cell)
                    out_rows.append({"source_id": int(src_id), "healpix_id": int(hid), "weight": weight})

        except Exception as e:
            # protect partition processing from crashing; log and continue
            logger.debug(f"skipping source {src_id} due to error: {e}")
            out_rows.append({"source_id": int(src_id) if src_id is not None else pd.NA, "healpix_id": pd.NA})
            total_count += 1
            dropped_count += 1
            continue

    # Log statistics for this partition
    if total_count > 0:
        drop_pct = 100.0 * dropped_count / total_count
        if dropped_count > 0:
            logger.info(f"Partition (lon_convention={lon_convention}): processed {total_count} geometries, "
                       f"dropped {dropped_count} ({drop_pct:.1f}%) total "
                       f"[pre-filter: {dropped_prefilter}, post-processing: {dropped_count - dropped_prefilter}]")
        else:
            logger.debug(f"Partition (lon_convention={lon_convention}): processed {total_count} geometries, no drops")
    
    if len(out_rows) == 0:
        return _pd.DataFrame(columns=["source_id", "healpix_id", "weight"])  # empty
    df_out = _pd.DataFrame(out_rows)
    # ensure types
    df_out["source_id"] = df_out["source_id"].astype(np.int64)
    df_out["healpix_id"] = df_out["healpix_id"].astype("UInt64")
    if 'weight' in df_out.columns:
        df_out["weight"] = df_out["weight"].astype(float)
    return df_out

# %% ../nbs/01_sidecar.ipynb 7
def add_psf_weights_to_sidecar(
    sidecar_df,
    src_geoms,
    cell_geoms,
    data_psf=None,
    cell_psf=None,
    combine_method='multiply',
    normalize=True
):
    """
    Add a 'weight' column to the sidecar DataFrame using PSF functions.
    - src_geoms: sequence of source geometries (indexed by source_id)
    - cell_geoms: dict or sequence mapping healpix_id to cell geometry
    """
    weights = []
    for row in sidecar_df.itertuples(index=False):
        src_id = row.source_id
        cell_id = row.healpix_id
        src_geom = src_geoms[src_id]
        cell_geom = cell_geoms[cell_id]
        w = compute_assignment_weight(src_geom, cell_geom, data_psf, cell_psf, combine_method)
        weights.append(w)
    sidecar_df = sidecar_df.copy()
    sidecar_df['weight'] = weights
    if normalize:
        sidecar_df = normalize_weights_per_cell(sidecar_df, cell_col='healpix_id', weight_col='weight')
    return sidecar_df

# %% ../nbs/01_sidecar.ipynb 8
def normalize_weights_per_cell(df, cell_col='healpix_id', weight_col='weight'):
    """Normalize weights so that sum of weights per cell is 1.0."""
    sums = df.groupby(cell_col)[weight_col].transform('sum')
    df[weight_col] = df[weight_col] / sums
    return df

# %% ../nbs/01_sidecar.ipynb 9
def compute_assignment_weight(
    src_geom,
    cell_geom,
    data_psf=None,
    cell_psf=None,
    combine_method='multiply',
    data_psf_sigma=None,
    cell_psf_sigma=None
):
    """
    Compute the assignment weight for a (source geometry, cell geometry) pair.
    - data_psf: callable or None
    - cell_psf: callable or None
    - combine_method: 'multiply', 'sum', 'min', 'max'
    """
    # For now, use centroid-to-centroid distance for polygons
    # (future: integrate over geometry or use rasterized PSF)
    src_centroid = src_geom.centroid
    cell_centroid = cell_geom.centroid
    dx = src_centroid.x - cell_centroid.x
    dy = src_centroid.y - cell_centroid.y
    w_data = data_psf(dx, dy) if data_psf else 1.0
    w_cell = cell_psf(dx, dy) if cell_psf else 1.0
    if combine_method == 'multiply':
        return w_data * w_cell
    elif combine_method == 'sum':
        return w_data + w_cell
    elif combine_method == 'min':
        return min(w_data, w_cell)
    elif combine_method == 'max':
        return max(w_data, w_cell)
    else:
        raise ValueError(f"Unknown combine_method: {combine_method}")

# %% ../nbs/01_sidecar.ipynb 10
def build_output_path(input_path: Path, mode: str, nside: int) -> Path:
    """Build output path for sidecar file based on input and parameters."""
    stem = input_path.stem
    # separate the key:values with - and _ between them 
    outname = f"{stem}.cell-healpix_assignment-{mode}_nside-{nside}_order-nested.parquet"
    return input_path.with_name(outname)


def write_sidecar_metadata(output_path: Path, input_path: Path, nside: int, mode: str, 
                           lon_convention: str, ncores: int, args) -> Path:
    """Write sidecar processing metadata to JSON file.
    
    Args:
        output_path: Path to the sidecar parquet file
        input_path: Path to the input file
        nside: HEALPix nside parameter
        mode: Assignment mode ('strict' or 'fuzzy')
        lon_convention: Longitude convention used
        ncores: Number of cores used
        args: Parsed command-line arguments
        
    Returns:
        Path to the written metadata file
    """
    from datetime import datetime, timezone
    from healpyxel.metadata import HEALPyxelxMetadata

    metadata = {
        'processing': {
            'stage': 'sidecar',
            'timestamp': datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
            'source_file': str(input_path.absolute()),
            'output_file': str(output_path.absolute())
        },
        'healpix': {
            'nside': nside,
            'mode': mode,
            'order': 'nested',
            'npix': 12 * nside ** 2
        },
        'coordinates': {
            'lon_convention': lon_convention,
            'lon_range': [0, 360] if lon_convention == '0_360' else [-180, 180],
            'lat_range': [-90, 90]
        },
        'processing_params': {
            'ncores': ncores,
            'coalesced': not args.no_coalesce,
            'data_psf': getattr(args, 'data_psf', None),
            'data_psf_sigma_level': getattr(args, 'data_psf_sigma_level', None),
            'cell_psf': getattr(args, 'cell_psf', None),
            'cell_psf_sigma_level': getattr(args, 'cell_psf_sigma_level', None),
            'psf_combine': getattr(args, 'psf_combine', None),
            'psf_normalize': not getattr(args, 'no_psf_normalize', False)
        }
    }

    return HEALPyxelxMetadata.write_json(metadata, output_path, validate=True)

# %% ../nbs/01_sidecar.ipynb 11
def validate_nside(nside: int) -> bool:
    """Validate that nside is a positive power of two."""
    return nside > 0 and (nside & (nside - 1)) == 0


def parse_arguments(argv=None):
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(description="Create HEALPix sidecar mapping source geometries to cells.")
    parser.add_argument("--input", "-i", required=True, help="Path to input GeoParquet file")
    parser.add_argument("--nside", "-n", type=int, nargs='+', default=None,
                        help="One or more HEALPix nside values (powers of 2). Required unless --geo-stats is used. Example: -n 64 128")
    parser.add_argument("--mode", "-m", choices=["strict", "fuzzy"], default="fuzzy",
                        help="Assignment mode: strict (single-cell only) or fuzzy (all cells touched)")
    parser.add_argument("--ncores", type=int, default=max(1, (os.cpu_count() or 2) - 1),
                        help="Number of cores to use for Dask workers (defaults to cpu_count-1)")
    parser.add_argument("--output_dir", "-o", default=None,
                        help="Directory to write the output file (defaults to same folder as input)")
    parser.add_argument("--no-coalesce", dest="no_coalesce", action="store_true",
                        help="Do not coalesce partitions into a single file; write partitioned parquet (default: coalesce to single file)")
    parser.add_argument("--lon-convention", type=str, default=None, choices=['0_360', 'minus_plus180'],
                        help="Longitude convention: '0_360' for [0,360) or 'minus_plus180' for [-180,180). Required unless --geo-stats is used. Can be used with --geo-stats to apply filtering.")
    parser.add_argument("--geo-stats", action="store_true",
                        help="Compute and display geographical statistics (lon/lat ranges, mean, std) before processing")
    parser.add_argument("--lon-col", type=str, default=None,
                        help="Longitude column name (if not specified, will auto-detect or extract from geometry)")
    parser.add_argument("--lat-col", type=str, default=None,
                        help="Latitude column name (if not specified, will auto-detect or extract from geometry)")
    parser.add_argument("--stats-sample-size", type=int, default=10000,
                        help="Number of rows to sample when extracting coordinates from geometry (default: 10000)")
    parser.add_argument("--loglevel", "-l", choices=["debug", "info", "warning", "error"], default="info",
                        help="Set logging level (default: info)")
    
    # Add PSF-related CLI arguments to an argparse parser
    parser.add_argument('--data-psf', type=str, default='none', choices=['none', 'gaussian'],
                        help='Data point spread function type (default: none)')
    parser.add_argument('--data-psf-sigma-level', type=float, default=2.0,
                        help='Sigma level for data PSF (default: 2.0)')
    parser.add_argument('--cell-psf', type=str, default='none', choices=['none', 'gaussian'],
                        help='Cell spread function type (default: none)')
    parser.add_argument('--cell-psf-sigma-level', type=float, default=2.0,
                        help='Sigma level for cell PSF (default: 2.0)')
    parser.add_argument('--psf-combine', type=str, default='multiply', choices=['multiply', 'sum', 'min', 'max'],
                        help='How to combine data and cell PSF weights (default: multiply)')
    parser.add_argument('--no-psf-normalize', action='store_true',
                        help='Disable normalization of weights per cell (default: normalize)')
 
    return parser.parse_args(argv)

# %% ../nbs/01_sidecar.ipynb 12
class PSF:
    """Base class for Point Spread Functions (PSF)."""
    def __init__(self):
        pass
    def __call__(self, dx, dy):
        raise NotImplementedError

class GaussianPSF(PSF):
    """2D Gaussian PSF centered at (0,0)."""
    def __init__(self, sigma=None):
        super().__init__()
        self.sigma = sigma  # If None, must be set by user or context
    def __call__(self, dx, dy):
        if self.sigma is None:
            raise ValueError("Sigma must be set for GaussianPSF.")
        r2 = dx**2 + dy**2
        return np.exp(-0.5 * r2 / (self.sigma**2))

PSF_REGISTRY = {
    'gaussian': GaussianPSF,
    'none': lambda *a, **k: 1.0,
}

def get_psf(psf_type, sigma=None):
    if psf_type == 'none':
        return lambda dx, dy: 1.0
    cls = PSF_REGISTRY.get(psf_type, None)
    if cls is None:
        raise ValueError(f"Unknown PSF type: {psf_type}")
    return cls(sigma=sigma)

# %% ../nbs/01_sidecar.ipynb 13
def write_partitioned_output(tasks, out_file: Path, nparts: int) -> int:
    """Write output as partitioned parquet files (one per partition).
    
    Returns:
        Total number of rows written
    """
    import dask
    
    out_part_dir = out_file.with_suffix('.parts')
    out_part_dir.mkdir(parents=True, exist_ok=True)
    logger.info(f"Writing partitioned parquet to {out_part_dir} (no coalesce requested, {nparts} partitions)")
    parts = dask.compute(*tasks)
    total_rows = 0
    files_written = 0
    for idx, df_part in enumerate(parts):
        if df_part is None or len(df_part) == 0:
            continue
        df_part = df_part.astype({"source_id": "int64", "healpix_id": "UInt64"})
        part_path = out_part_dir / f"part-{idx:06d}.parquet"
        try:
            df_part.to_parquet(str(part_path), engine="pyarrow", index=False)
            total_rows += len(df_part)
            files_written += 1
        except Exception:
            # fallback: write via pyarrow Table to ensure metadata compatibility
            try:
                import pyarrow as pa
                table = pa.Table.from_pandas(df_part, preserve_index=False)
                import pyarrow.parquet as pq
                pq.write_table(table, str(part_path))
                total_rows += len(df_part)
                files_written += 1
            except Exception:
                logger.exception("Failed to write partitioned parquet for %s", part_path)
    logger.info(f"Wrote partitioned parquet to {out_part_dir}: {files_written} files, {total_rows} total rows")
    return total_rows

# %% ../nbs/01_sidecar.ipynb 14
def write_coalesced_output(tasks, out_file: Path, nside: int, mode: str, ncores: int, nparts: int) -> int:
    """Write output as a single coalesced parquet file with incremental batching.

    Returns:
        Total number of rows written
    """
    import dask
    import pyarrow as pa
    import pyarrow.parquet as pq

    logger.info(f"Computing {nparts} partitions and writing single parquet file to {out_file}")

    if nparts == 0:
        # nothing to write; write empty file with explicit schema and metadata
        schema = pa.schema([
            ("source_id", pa.int64()),
            ("healpix_id", pa.uint64()),
        ])
        pq_meta = {"nside": str(nside), "mode": mode, "order": "nested"}
        schema = schema.with_metadata({k: v.encode() for k, v in pq_meta.items()})
        empty_table = pa.Table.from_pandas(
            pd.DataFrame(columns=["source_id", "healpix_id"]).astype({"source_id": "int64", "healpix_id": "UInt64"}),
            schema=schema, preserve_index=False
        )
        pq.write_table(empty_table, str(out_file))
        logger.info(f"Wrote empty output {out_file}")
        return 0

    # --- Patch: Dynamically detect if 'weight' column is present in the first non-empty batch ---
    batch_size = max(1, ncores)
    schema = None
    pq_meta = {"nside": str(nside), "mode": mode, "order": "nested"}
    first_batch_has_weight = False

    # Find the first non-empty batch to determine schema
    for i in range(0, nparts, batch_size):
        batch = tasks[i : i + batch_size]
        import dask
        res = dask.compute(*batch)
        non_empty = [r for r in res if (r is not None and len(r) > 0)]
        if not non_empty:
            continue
        df_batch = pd.concat(non_empty, ignore_index=True)
        # Check for weight column
        if "weight" in df_batch.columns:
            first_batch_has_weight = True
        break  # Only need to check the first non-empty batch

    # Build schema accordingly
    if first_batch_has_weight:
        schema = pa.schema([
            ("source_id", pa.int64()),
            ("healpix_id", pa.uint64()),
            ("weight", pa.float64()),
        ])
    else:
        schema = pa.schema([
            ("source_id", pa.int64()),
            ("healpix_id", pa.uint64()),
        ])
    schema = schema.with_metadata({k: v.encode() for k, v in pq_meta.items()})

    writer = None
    total_rows_written = 0
    total_batches_processed = 0

    # Per-nside progress bar that advances as Dask computes partition batches.
    pbar = tqdm(total=nparts, desc=f"nside={nside}", unit="part", position=0, leave=True)
    try:
        for i in range(0, nparts, batch_size):
            batch = tasks[i : i + batch_size]
            batch_start = i + 1
            batch_end = min(i + batch_size, nparts)
            pbar.set_description(f"nside={nside} [{batch_start}-{batch_end}/{nparts}]")

            # compute this batch with progress tracking if available
            if DASK_PROGRESS_AVAILABLE and logger.level <= logging.INFO:
                with DaskProgressBar():
                    res = dask.compute(*batch)
            else:
                res = dask.compute(*batch)
            # res is a tuple of pandas DataFrames
            non_empty = [r for r in res if (r is not None and len(r) > 0)]
            if not non_empty:
                pbar.update(len(batch))
                continue
            df_batch = pd.concat(non_empty, ignore_index=True)
            # ensure dtypes
            df_batch = df_batch.astype({"source_id": "int64", "healpix_id": "UInt64"})
            # If schema has weight but batch does not, add column
            if first_batch_has_weight and "weight" not in df_batch.columns:
                df_batch["weight"] = np.nan
            # create arrow table with enforced schema/metadata
            table = pa.Table.from_pandas(df_batch, schema=schema, preserve_index=False)
            batch_rows = len(df_batch)
            total_rows_written += batch_rows
            total_batches_processed += 1

            if writer is None:
                # overwrite existing single file if present
                if out_file.exists():
                    try:
                        out_file.unlink()
                    except Exception:
                        logger.warning(f"Could not remove existing output file {out_file}")
                writer = pq.ParquetWriter(str(out_file), schema)
            writer.write_table(table)

            # update progress with statistics
            pbar.set_postfix({
                'rows': total_rows_written,
                'batch_rows': batch_rows
            })
            pbar.update(len(batch))
    finally:
        pbar.close()

    # close the writer if we created one
    if writer is not None:
        writer.close()
        logger.info(f"Wrote single parquet file: {out_file} ({total_rows_written} rows, {total_batches_processed} batches)")
    else:
        # if nothing was written, write empty file with schema
        logger.info(f"Wrote empty output {out_file}")
        empty_table = pa.Table.from_pandas(
            pd.DataFrame(columns=["source_id", "healpix_id"]).astype({"source_id": "int64", "healpix_id": "UInt64"}),
            schema=schema, preserve_index=False
        )
        pq.write_table(empty_table, str(out_file))

    return total_rows_written

# %% ../nbs/01_sidecar.ipynb 15
def main(argv=None):
    """Main entry point for HEALPix sidecar generation."""
    args = parse_arguments(argv)

    input_path = Path(args.input)
    if not input_path.exists():
        logger.error(f"Input path does not exist: {input_path}")
        sys.exit(2)

    # Validate required parameters based on mode
    if not args.geo_stats:
        # When not using --geo-stats, require --nside and --lon-convention
        if args.nside is None:
            logger.error("--nside is required when not using --geo-stats")
            sys.exit(2)
        if args.lon_convention is None:
            logger.error("--lon-convention is required when not using --geo-stats")
            sys.exit(2)

    # validate nsides (only if provided)
    nsides = args.nside
    if nsides is not None:
        for n in nsides:
            if not validate_nside(n):
                logger.error("nside must be a positive power of two: invalid value %s", n)
                sys.exit(2)

    # configure logging level from CLI
    level_map = {
        'debug': logging.DEBUG,
        'info': logging.INFO,
        'warning': logging.WARNING,
        'error': logging.ERROR,
    }
    root_level = level_map.get(args.loglevel, logging.INFO)
    logging.getLogger().setLevel(root_level)
    logger.setLevel(root_level)

    # Compute and display geo-statistics if requested (and exit)
    if args.geo_stats:
        if args.lon_convention:
            logger.info(f"Computing geographical statistics with filtering (--lon-convention {args.lon_convention})...")
        else:
            logger.info("Computing geographical statistics (raw data, no filtering)...")
        try:
            stats = compute_geo_statistics(
                input_path, 
                lon_col=args.lon_col,
                lat_col=args.lat_col,
                sample_size=args.stats_sample_size,
                lon_convention=args.lon_convention
            )
            
            if stats:
                formatted_stats = format_geo_statistics(stats)
                print("\n" + formatted_stats + "\n")
                
                # Save statistics to JSON file
                stats_output = input_path.with_suffix('.geo_stats.json')
                try:
                    import json
                    with open(stats_output, 'w') as f:
                        json.dump(stats, f, indent=2)
                    logger.info(f"Saved geo-statistics to {stats_output}")
                except Exception as e:
                    logger.debug(f"Could not save statistics to JSON: {e}")
                
                # Exit successfully after showing statistics
                logger.info("Geo-statistics complete. Exiting (use without --geo-stats to process data).")
                sys.exit(0)
            else:
                logger.error("Could not compute geo-statistics")
                sys.exit(1)
        except Exception as e:
            logger.error(f"Geo-statistics computation failed: {e}")
            if logger.level <= logging.DEBUG:
                import traceback
                traceback.print_exc()
            sys.exit(1)

    logger.info(f"Reading input lazily from {input_path}; ncores={args.ncores}; mode={args.mode}; nsides={nsides}")
    logger.info(f"Longitude convention: {args.lon_convention}")

    # read lazily with dask_geopandas
    # let dask decide partitions but hint with npartitions based on ncores
    try:
        ddf = dg.read_parquet(str(input_path))
    except Exception:
        # try forcing to use dask read with explicit npartitions
        ddf = dg.read_parquet(str(input_path), npartitions=args.ncores)

    # We will compute explicit, global `source_id` values (0..N-1) that match
    # `gdf.reset_index()` by computing per-partition offsets and passing them to
    # `process_partition` as `base_index`.

    # apply partition-level processing
    # map_partitions will return a dask DataFrame; ensure meta is correct
    # Dask/meta: use pandas-recognized dtype strings to avoid version-specific dtype issues
    meta = pd.DataFrame({"source_id": pd.Series(dtype="int64"), "healpix_id": pd.Series(dtype="UInt64")})

    logger.info("Starting partitioned HEALPix assignment (this may take time)")

    # Loop over requested nsides without reloading the input (show overall progress)
    for nside in tqdm(nsides, desc="nsides", unit="nside"):
        logger.info("Processing nside=%s", nside)
        # We'll use delayed partitions so we can compute global source_id offsets
        import dask
        delayed_partitions = ddf.to_delayed()
        # compute partition lengths to derive base offsets for source_id
        try:
            part_lengths = ddf.map_partitions(lambda df: len(df)).compute().tolist()
        except Exception:
            # fallback: compute lengths by materializing delayed partitions (may be slower)
            part_lengths = [int(dask.compute(dask.delayed(lambda df: len(df))(p))[0]) for p in delayed_partitions]

        if len(part_lengths) != len(delayed_partitions):
            # defensive: if mismatch, fallback to equal-sized offsets (best-effort)
            logger.warning("Partition count mismatch; falling back to equal offsets")
            part_lengths = [len(delayed_partitions)] * len(delayed_partitions)

        offsets = np.concatenate(([0], np.cumsum(part_lengths)[:-1]))
        
        data_psf = None
        cell_psf = None        
        if args.cell_psf != 'none':
            logger.info(f"Using cell PSF: {args.cell_psf} (sigma_level={args.cell_psf_sigma_level})")
            # Set sigma to a reasonable value, e.g., 1.0 or based on nside
            cell_psf = get_psf(args.cell_psf, sigma=args.cell_psf_sigma_level)
        if args.data_psf != 'none':
            logger.info(f"Using data PSF: {args.data_psf} (sigma_level={args.data_psf_sigma_level})")
            data_psf = get_psf(args.data_psf, sigma=args.cell_psf_sigma_level)

        # build delayed tasks that pass the base_index per partition so source_id
        # corresponds to global row numbering (like geopandas.reset_index())
        tasks = [dask.delayed(process_partition)(
                    part, nside, args.mode, int(offsets[i]), args.lon_convention,
                    data_psf=data_psf, cell_psf=cell_psf, combine_method=args.psf_combine
                    )
                    for i, part in enumerate(delayed_partitions)
                ]
        nparts = len(tasks)

        # prepare output path (single file target by default)
        if args.output_dir:
            out_dir = Path(args.output_dir)
            out_dir.mkdir(parents=True, exist_ok=True)
        else:
            out_dir = input_path.parent

        out_file = out_dir / build_output_path(input_path, args.mode, nside).name

        # If user requested partitioned (no coalesce) output, compute each delayed
        # partition and write a separate parquet file per partition while preserving
        # global `source_id` numbering.
        if args.no_coalesce:
            total_rows = write_partitioned_output(tasks, out_file, nparts)
            # Write metadata
            metadata_path = write_sidecar_metadata(
                out_file, input_path, nside, args.mode, 
                args.lon_convention, args.ncores, args
            )
            logger.info(f"Wrote metadata to {metadata_path}")
            continue

        # compute partitions in batches and write incrementally using pyarrow ParquetWriter
        try:
            import dask
            import pyarrow as pa
            import pyarrow.parquet as pq
        except Exception as e:
            logger.error("pyarrow and dask are required for coalescing partitions to a single file: %s", e)
            logger.info("Falling back to writing partitioned parquet folder instead")
            out_part_dir = str(out_file.with_suffix('.parts'))
            result.to_parquet(out_part_dir, engine="pyarrow", write_index=False)
            logger.info("Wrote partitioned parquet to %s", out_part_dir)
            continue

        total_rows = write_coalesced_output(tasks, out_file, nside, args.mode, args.ncores, nparts)
        
        # Write metadata
        metadata_path = write_sidecar_metadata(
            out_file, input_path, nside, args.mode, 
            args.lon_convention, args.ncores, args
        )
        logger.info(f"Wrote metadata to {metadata_path}")

    # end loop over nsides


# CLI entry point (use via command line or import main() function)

# %% ../nbs/01_sidecar.ipynb 16
def get_healpix_cell_geometry(healpix_id, nside, nest=True):
    """
    Return a shapely Polygon for the given HEALPix cell.
    Uses healpy boundaries (in degrees, lon/lat).
    """
    import healpy as hp
    from shapely.geometry import Polygon

    # Get the boundary vertices of the cell (returns theta, phi in radians)
    vertices = hp.boundaries(nside, healpix_id, step=1, nest=nest)  # shape (2, N)
    # Convert to lon/lat in degrees
    theta = vertices[0]  # colatitude in radians
    phi = vertices[1]    # longitude in radians
    lats = 90.0 - np.degrees(theta)
    lons = np.degrees(phi)
    # Build polygon (healpy returns vertices in order)
    coords = list(zip(lons, lats))
    return Polygon(coords)
