# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_finalize.ipynb.

# %% auto 0
__all__ = ['logger', 'finalize_statistics', 'densify_healpix_map', 'export_to_geotiff', 'main']

# %% ../nbs/04_finalize.ipynb 1
"""healpix_finalize.py

Convert accumulator state to final HEALPix maps with statistics.

Takes the streaming state file produced by healpix_accumulator.py and computes
final statistics (mean, std, percentiles) for each HEALPix cell. Optionally
densifies the output to create a complete HEALPix grid.

Features:
- Compute mean, std, min, max from streaming statistics
- Approximate percentiles from T-Digest (if available)
- Optional densification to full HEALPix grid
- Quality control via minimum observation count
- Export to parquet, GeoTIFF, or other formats

Requirements:
- pandas, numpy, pyarrow (core)
- healpy (for densification, optional)
- rasterio (for GeoTIFF export, optional)

Usage example:
  # Basic finalization
  healpix_finalize \
    --state state/state_v030.parquet \
    --output products/month01_mosaic.parquet \
    --min-count 5

  # With percentiles and densification
  healpix_finalize \
    --state state/state_v030.parquet \
    --output products/month01_mosaic_dense.parquet \
    --percentiles 25 50 75 \
    --min-count 10 \
    --densify --nside 512
"""

from pathlib import Path
from typing import Optional, List, Dict, Any, Tuple
import logging
import argparse
import json
from datetime import datetime, timezone
import sys

import numpy as np
import pandas as pd

from .metadata import HEALPyxelxMetadata, FileType

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

# Import from accumulator
try:
    from healpyxel.accumulator import load_state, CellAccumulator, TDIGEST_AVAILABLE
except ImportError:
    try:
        # Fallback to standalone script
        from healpix_accumulator import load_state, CellAccumulator, TDIGEST_AVAILABLE
    except ImportError:
        logger.warning("Could not import from healpix_accumulator, using local implementation")
        load_state = None
        CellAccumulator = None
        TDIGEST_AVAILABLE = False

try:
    import healpy as hp
    HEALPY_AVAILABLE = True
except ImportError:
    HEALPY_AVAILABLE = False
    hp = None

try:
    from tqdm.auto import tqdm
    TQDM_AVAILABLE = True
except ImportError:
    TQDM_AVAILABLE = False
    tqdm = None

# %% ../nbs/04_finalize.ipynb 2
def finalize_statistics(
    state: Dict[int, CellAccumulator],
    percentiles: Optional[List[float]] = None,
    min_count: int = 1,
) -> pd.DataFrame:
    """
    Convert accumulator state to final statistics DataFrame.
    
    Args:
        state: Accumulator state dictionary {healpix_id: CellAccumulator}
        percentiles: List of percentiles to compute (e.g., [25, 50, 75])
        min_count: Minimum observations required per cell (cells below this are NaN)
        
    Returns:
        DataFrame indexed by healpix_id with statistics columns:
        - {col}_n: observation count
        - {col}_mean: mean value
        - {col}_std: standard deviation
        - {col}_min: minimum value
        - {col}_max: maximum value
        - {col}_p{N}: percentile (if T-Digest available)
    """
    if percentiles is None:
        percentiles = []
    
    logger.info(f"Finalizing statistics for {len(state)} cells")
    logger.info(f"Minimum observation count: {min_count}")
    
    rows = []
    iterator = state.items() if not TQDM_AVAILABLE else tqdm(state.items(), desc="Computing statistics")
    
    for hp_id, acc in iterator:
        row = {'healpix_id': int(hp_id)}
        
        # Get all columns processed
        columns = list(acc.stats_by_column.keys())
        
        for col in columns:
            stats = acc.stats_by_column[col]
            
            # Always include count
            row[f'{col}_n'] = int(stats.n)
            
            # Skip cells with insufficient data
            if stats.n < min_count:
                row[f'{col}_mean'] = float('nan')
                row[f'{col}_std'] = float('nan')
                row[f'{col}_min'] = float('nan')
                row[f'{col}_max'] = float('nan')
                
                for p in percentiles:
                    row[f'{col}_p{int(p)}'] = float('nan')
                continue
            
            # Basic statistics
            row[f'{col}_mean'] = stats.mean
            row[f'{col}_std'] = stats.std
            row[f'{col}_min'] = stats.min_val if np.isfinite(stats.min_val) else float('nan')
            row[f'{col}_max'] = stats.max_val if np.isfinite(stats.max_val) else float('nan')
            
            # Percentiles from T-Digest
            if hasattr(acc, 'tdigests') and col in acc.tdigests:
                digest = acc.tdigests[col]
                for p in percentiles:
                    try:
                        value = digest.percentile(p)
                        row[f'{col}_p{int(p)}'] = float(value)
                    except Exception as e:
                        logger.warning(f"Failed to compute p{p} for {col} in cell {hp_id}: {e}")
                        row[f'{col}_p{int(p)}'] = float('nan')
            else:
                # No T-Digest available, set to NaN
                for p in percentiles:
                    row[f'{col}_p{int(p)}'] = float('nan')
        
        rows.append(row)
    
    df = pd.DataFrame(rows).set_index('healpix_id').sort_index()
    
    logger.info(f"✓ Finalized {len(df)} cells")
    logger.info(f"  Columns: {list(df.columns)}")
    
    # Report coverage statistics
    for col in columns:
        if f'{col}_n' in df.columns:
            valid_cells = (df[f'{col}_n'] >= min_count).sum()
            total_obs = df[f'{col}_n'].sum()
            logger.info(f"  {col}: {valid_cells} valid cells, {int(total_obs):,} total observations")
    
    return df

# %% ../nbs/04_finalize.ipynb 3
def densify_healpix_map(
    sparse_df: pd.DataFrame,
    nside: int,
    fill_value: float = np.nan
) -> pd.DataFrame:
    """
    Create a complete HEALPix grid by filling empty cells with fill_value.
    
    Args:
        sparse_df: DataFrame with healpix_id index (sparse)
        nside: HEALPix nside parameter
        fill_value: Value for empty cells (default: NaN)
        
    Returns:
        Dense DataFrame with all 12*nside**2 cells
    """
    if not HEALPY_AVAILABLE:
        logger.error("healpy required for densification (pip install healpy)")
        raise ImportError("healpy not available")
    
    n_pixels = hp.nside2npix(nside)
    logger.info(f"Densifying to full grid (nside={nside}, {n_pixels} cells)")
    
    # Create full index
    full_index = pd.Index(range(n_pixels), name='healpix_id')
    
    # Reindex with fill_value
    dense_df = sparse_df.reindex(full_index, fill_value=fill_value)
    
    logger.info(f"✓ Densified: {len(sparse_df)} → {len(dense_df)} cells")
    
    return dense_df

# %% ../nbs/04_finalize.ipynb 4
def export_to_geotiff(
    df: pd.DataFrame,
    column: str,
    output_path: Path,
    nside: int,
    crs: str = 'IAU:19900',  # Mercury IAU CRS
):
    """
    Export a column to GeoTIFF format (requires rasterio and healpy).
    
    Note: This creates an equirectangular projection from HEALPix data.
    """
    try:
        import rasterio
        from rasterio.transform import from_bounds
    except ImportError:
        logger.error("rasterio required for GeoTIFF export (pip install rasterio)")
        raise
    
    if not HEALPY_AVAILABLE:
        logger.error("healpy required for GeoTIFF export (pip install healpy)")
        raise ImportError("healpy not available")
    
    logger.info(f"Exporting {column} to GeoTIFF: {output_path}")
    
    # Convert to healpy array
    n_pixels = hp.nside2npix(nside)
    healpix_map = np.full(n_pixels, np.nan, dtype=np.float32)
    
    for hp_id, value in df[column].items():
        if 0 <= hp_id < n_pixels:
            healpix_map[hp_id] = value
    
    # Convert HEALPix to equirectangular grid
    width, height = 1440, 720  # 0.25 deg resolution
    lon = np.linspace(-180, 180, width)
    lat = np.linspace(90, -90, height)
    lon_grid, lat_grid = np.meshgrid(lon, lat)
    
    # HEALPix uses colatitude (0 at north pole)
    theta = np.radians(90 - lat_grid)
    phi = np.radians(lon_grid)
    
    # Query HEALPix values
    pixels = hp.ang2pix(nside, theta, phi, nest=True)
    grid = healpix_map[pixels]
    
    # Write GeoTIFF
    transform = from_bounds(-180, -90, 180, 90, width, height)
    
    with rasterio.open(
        output_path,
        'w',
        driver='GTiff',
        height=height,
        width=width,
        count=1,
        dtype=grid.dtype,
        crs=crs,
        transform=transform,
        compress='deflate',
        nodata=np.nan,
    ) as dst:
        dst.write(grid, 1)
    
    logger.info(f"✓ Exported GeoTIFF ({width}x{height})")

# %% ../nbs/04_finalize.ipynb 5
def _normalize_load_state_result(result) -> Tuple[Dict[int, CellAccumulator], Optional[HEALPyxelxMetadata]]:
    """Normalize load_state outputs across versions."""
    if isinstance(result, tuple) and len(result) == 2:
        return result
    return result, None


def _in_ipython_kernel() -> bool:
    """Return True when running inside an IPython kernel (notebook)."""
    try:
        from IPython import get_ipython
        ip = get_ipython()
        return ip is not None and 'IPKernelApp' in ip.config
    except Exception:
        return False


def main(argv=None):
    parser = argparse.ArgumentParser(
        description="Finalize accumulator state to statistical HEALPix maps",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Basic finalization
  python healpix_finalize.py --state state_v030.parquet --output month01_mosaic.parquet
  
  # With percentiles and minimum count
  python healpix_finalize.py --state state_v030.parquet --output month01.parquet \\
    --percentiles 25 50 75 --min-count 10
  
  # Densify to full grid
  python healpix_finalize.py --state state_v030.parquet --output month01_dense.parquet \\
    --densify --nside 512
  
  # Export specific column to GeoTIFF
  python healpix_finalize.py --state state_v030.parquet --output month01.parquet \\
    --export-tiff r750_mean output_r750.tif --nside 512
"""
    )
    
    # Input/output
    parser.add_argument('-s', '--state', required=True, type=Path,
                       help='Input state file from healpix_accumulator.py')
    parser.add_argument('-o', '--output', required=True, type=Path,
                       help='Output parquet file with statistics')
    
    # Statistics options
    parser.add_argument('-p', '--percentiles', type=float, nargs='+',
                       default=[25, 50, 75],
                       help='Percentiles to compute (default: 25 50 75)')
    parser.add_argument('--min-count', type=int, default=1,
                       help='Minimum observations per cell (default: 1)')
    
    # Densification
    parser.add_argument('--densify', action='store_true',
                       help='Create full HEALPix grid (fill empty cells with NaN)')
    parser.add_argument('--nside', type=int,
                       help='nside for densification (auto-detected from metadata if not specified)')
    
    # Export options
    parser.add_argument('--export-tiff', nargs=2, metavar=('COLUMN', 'OUTPUT'),
                       help='Export a column to GeoTIFF (requires rasterio)')
    parser.add_argument('--crs', default='IAU:19900',
                       help='CRS for GeoTIFF export (default: IAU:19900 for Mercury)')
    
    # Logging
    parser.add_argument('-v', '--verbose', action='store_true',
                       help='Enable verbose debug logging')
    parser.add_argument('-q', '--quiet', action='store_true',
                       help='Suppress progress bars and non-critical messages')
    
    args = parser.parse_args(argv)
    
    # Configure logging
    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)
    elif args.quiet:
        logging.getLogger().setLevel(logging.WARNING)
    
    # Validate inputs
    if not args.state.exists():
        logger.error(f"State file not found: {args.state}")
        return 1
    
    # Load and validate HEALPix metadata from state
    try:
        meta_in = HEALPyxelxMetadata.from_parquet(args.state)
        logger.info(f"State metadata: nside={meta_in.nside}, mode={meta_in.mode}, order={meta_in.order}")
    except ValueError as e:
        logger.error(f"State file missing HEALPix metadata: {e}")
        return 1
    
    meta_out = HEALPyxelxMetadata(
        nside=meta_in.nside,
        order=meta_in.order,
        npix=meta_in.npix,
        mode=meta_in.mode,
        lon_convention=meta_in.lon_convention,
        file_type=FileType.FINALIZE,
    )
    
    if args.densify or args.export_tiff:
        if not args.nside:
            args.nside = meta_in.nside
            logger.info(f"Auto-detected nside={args.nside} from metadata")
    
    # Load state
    if load_state is None:
        logger.error("Could not import load_state from healpix_accumulator.py")
        logger.error("Ensure healpix_accumulator.py is in the same directory")
        return 1
    
    logger.info(f"Loading state from {args.state}")
    try:
        state_result = load_state(args.state, use_tdigest=True)
        state, _ = _normalize_load_state_result(state_result)
    except Exception as e:
        logger.error(f"Failed to load state: {e}")
        return 1
    
    # Check for T-Digest data
    has_tdigest = any(
        hasattr(acc, 'tdigests') and len(acc.tdigests) > 0
        for acc in state.values()
    )
    
    if args.percentiles and not has_tdigest:
        logger.warning("T-Digest data not found in state, percentiles will be NaN")
        logger.warning("Use healpix_accumulator.py without --no-tdigest to enable percentiles")
    
    # Finalize statistics
    logger.info("Computing statistics...")
    df = finalize_statistics(
        state=state,
        percentiles=args.percentiles if has_tdigest else [],
        min_count=args.min_count,
    )
    
    # Densify if requested
    if args.densify:
        if not HEALPY_AVAILABLE:
            logger.error("healpy required for densification (pip install healpy)")
            return 1
        df = densify_healpix_map(df, nside=args.nside)
    
    # Save output
    logger.info(f"Saving to {args.output}")
    args.output.parent.mkdir(parents=True, exist_ok=True)
    df.to_parquet(
        args.output,
        engine='pyarrow',
        compression='snappy',
        schema_metadata=meta_out.to_parquet_metadata(),
    )
    
    # Save metadata
    full_metadata = {
        'processing': {
            'stage': 'finalize',
            'timestamp': datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
            'source_state': str(args.state),
            'output_file': str(args.output),
            'n_cells': len(df),
            'percentiles': args.percentiles if has_tdigest else [],
            'min_count': args.min_count,
            'densified': args.densify,
            'export_tiff': args.export_tiff,
        },
        'healpix': {
            'nside': meta_out.nside,
            'mode': meta_out.mode,
            'order': meta_out.order,
            'npix': meta_out.npix,
        },
        'coordinates': {
            'lon_convention': meta_out.lon_convention,
            'lon_range': [0, 360] if meta_out.lon_convention == '0_360' else [-180, 180],
            'lat_range': [-90, 90],
        },
    }
    
    HEALPyxelxMetadata.write_json(full_metadata, args.output, validate=True)
    logger.info(f"✓ Saved metadata to {args.output.with_suffix('.meta.json')}")
    
    # Export to GeoTIFF if requested
    if args.export_tiff:
        column, tiff_path = args.export_tiff
        if column not in df.columns:
            logger.error(f"Column {column} not found in output")
            logger.error(f"Available columns: {list(df.columns)}")
            return 1
        
        try:
            export_to_geotiff(
                df=df,
                column=column,
                output_path=Path(tiff_path),
                nside=args.nside,
                crs=args.crs,
            )
        except Exception as e:
            logger.error(f"GeoTIFF export failed: {e}")
            return 1
    
    logger.info("✓ Finalization complete!")
    logger.info(f"  Output: {args.output}")
    logger.info(f"  Cells: {len(df)}")
    
    return 0

# %% ../nbs/04_finalize.ipynb 6
#| eval: false
if __name__ == '__main__':
    sys.exit(main())
    if _in_ipython_kernel():
        logger.info("Notebook context detected; skipping CLI entrypoint.")
    else:
        sys.exit(main())
