# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00b_metadata.ipynb.

# %% auto 0
__all__ = ['FileType', 'HEALPyxelxMetadata', 'validate_sidecar', 'validate_aggregate', 'save_with_metadata',
           'load_with_validation', 'write_metadata_with_file', 'read_metadata_from_file', 'validate_file_type',
           'load_cli_aggregate']

# %% ../nbs/00b_metadata.ipynb 2
from dataclasses import dataclass, asdict
from typing import Literal, Optional, Union
import pandas as pd
import pyarrow.parquet as pq
import healpy as hp
import numpy as np
import json
from pathlib import Path

# %% ../nbs/00b_metadata.ipynb 3
from enum import Enum

class FileType(str, Enum):
    """Pipeline file type discriminator.
    
    Used to identify which stage of the HEALPix pipeline produced a parquet file.
    """
    SIDECAR = 'sidecar'
    AGGREGATE = 'aggregate'
    ACCUMULATOR = 'accumulator'
    FINALIZE = 'finalize'
    GEOSPATIAL = 'geospatial'

# %% ../nbs/00b_metadata.ipynb 5
@dataclass
class HEALPyxelxMetadata:
    """Validated HEALPix grid metadata with embedded parquet I/O.
    
    All aggregate/sidecar parquet files should embed this in their schema
    to ensure metadata/data consistency.
    
    Attributes:
        nside: HEALPix resolution parameter (must be power of 2)
        order: Pixel ordering scheme ('nested' or 'ring')
        npix: Total number of pixels (auto-computed from nside if None)
        mode: Cell assignment mode ('strict' or 'fuzzy')
        lon_convention: Longitude convention ('0_360' or '-180_180')
        file_type: Pipeline stage that produced this file (optional)
    
    Examples:
        >>> meta = HEALPyxelxMetadata(nside=32, order='nested')
        >>> meta.npix
        12288
        >>> meta.nest
        True
    """
    nside: int
    order: Literal['nested', 'ring'] = 'nested'
    npix: Optional[int] = None
    mode: Literal['strict', 'fuzzy'] = 'fuzzy'
    lon_convention: Literal['0_360', '-180_180'] = '0_360'
    file_type: Optional[FileType] = None
    
    def __post_init__(self):
        """Validate and compute derived fields."""
        from healpyxel.core import validate_nside
        self.nside = validate_nside(self.nside)
        
        if self.npix is None:
            self.npix = hp.nside2npix(self.nside)
        
        expected_npix = hp.nside2npix(self.nside)
        assert self.npix == expected_npix, (
            f"npix={self.npix} inconsistent with nside={self.nside} "
            f"(expected {expected_npix})"
        )
        
        # Convert string file_type to enum if needed
        if isinstance(self.file_type, str):
            self.file_type = FileType(self.file_type)
    
    @classmethod
    def from_dict(cls, d: dict) -> 'HEALPyxelxMetadata':
        """Parse from nested or flat metadata dict.
        
        Supports both nested structure (from JSON sidecars) and flat dicts.
        
        Args:
            d: Dictionary with HEALPix parameters
        
        Returns:
            Validated HEALPyxelxMetadata instance
        """
        if 'healpix' in d or 'sidecar_metadata' in d:
            # Nested structure: {'sidecar_metadata': {'healpix': {...}}}
            hp_dict = d.get('sidecar_metadata', {}).get('healpix', d.get('healpix', {}))
            coord_dict = d.get('sidecar_metadata', {}).get('coordinates', {})
            lon_convention = coord_dict.get('lon_convention', d.get('lon_convention', '0_360'))
        else:
            # Flat structure
            hp_dict = d
            lon_convention = d.get('lon_convention', '0_360')
        
        # Try multiple keys for file_type (nested vs flat)
        file_type = d.get('file_type') or d.get('processing', {}).get('stage')
        if isinstance(file_type, str):
            try:
                file_type = FileType(file_type)
            except ValueError:
                file_type = None
        
        return cls(
            nside=int(hp_dict['nside']),
            order=hp_dict.get('order', 'nested'),
            npix=hp_dict.get('npix'),
            mode=hp_dict.get('mode', 'fuzzy'),
            lon_convention=lon_convention,
            file_type=file_type,
        )
    
    @classmethod
    def from_parquet(cls, path: Union[str, Path]) -> 'HEALPyxelxMetadata':
        """Extract metadata from parquet schema or companion JSON.
        
        Tries embedded parquet metadata first, falls back to .meta.json file.
        
        Args:
            path: Path to parquet file
        
        Returns:
            Parsed HEALPyxelxMetadata
        
        Raises:
            ValueError: If no metadata found
        """
        path = Path(path)
        
        # Try embedded parquet metadata first
        try:
            pf = pq.ParquetFile(path)
            schema_meta = pf.schema_arrow.metadata
            if schema_meta and b'healpix_metadata' in schema_meta:
                meta_json = schema_meta[b'healpix_metadata'].decode()
                return cls.from_dict(json.loads(meta_json))
        except Exception:
            pass
        
        # Fallback to companion .meta.json
        json_path = path.with_suffix('.meta.json')
        if json_path.exists():
            with open(json_path) as f:
                return cls.from_dict(json.load(f))
        
        raise ValueError(f"No HEALPix metadata found for {path}")
    
    def to_dict(self) -> dict:
        """Serialize to dictionary, converting Enum to string.
        
        Returns:
            Dictionary with all metadata fields
        """
        d = asdict(self)
        if self.file_type is not None:
            d['file_type'] = self.file_type.value
        return d

    @classmethod
    def write_json(
        cls,
        metadata: dict,
        output_file: Union[str, Path],
        validate: bool = True
    ) -> Path:
        """Validate metadata (optional), sanitize for JSON, and write .meta.json.
        
        Args:
            metadata: Metadata dict to serialize
            output_file: Full path to data file being described
            validate: If True, run HEALPyxelxMetadata.from_dict() sanity checks
        
        Returns:
            Path to written .meta.json file
        """
        import json as _json
        import datetime as _dt
        from pathlib import Path as _Path
        import numpy as _np

        def _jsonify(obj):
            if isinstance(obj, _Path):
                return str(obj)
            if isinstance(obj, (_dt.datetime, _dt.date)):
                return obj.isoformat()
            if isinstance(obj, _np.generic):
                return obj.item()
            if isinstance(obj, dict):
                return {str(k): _jsonify(v) for k, v in obj.items()}
            if isinstance(obj, (list, tuple)):
                return [_jsonify(v) for v in obj]
            return obj

        if validate:
            cls.from_dict(metadata)

        output_file = _Path(output_file)
        metadata_path = output_file.with_suffix('.meta.json')

        with open(metadata_path, 'w') as f:
            _json.dump(_jsonify(metadata), f, indent=2)

        return metadata_path
    
    def to_parquet_metadata(self) -> dict:
        """Serialize for embedding in parquet schema.
        
        Returns:
            Dictionary suitable for pyarrow schema metadata
        """
        return {b'healpix_metadata': json.dumps(self.to_dict()).encode()}
    
    @property
    def nest(self) -> bool:
        """Convenience flag for healpy calls."""
        return self.order == 'nested'

# %% ../nbs/00b_metadata.ipynb 7
def validate_sidecar(
    df: pd.DataFrame,
    meta: HEALPyxelxMetadata,
    strict: bool = True
) -> dict:
    """Validate sidecar structure and return diagnostic metrics.
    
    Checks:
    - Required columns (source_id, healpix_id)
    - healpix_id range [0, npix)
    - Unique cells ≤ npix
    - Weight validity [0, 1] for fuzzy mode
    - Weight sum ≈ 1.0 per source for fuzzy mode
    - 1:1 mapping for strict mode
    
    Args:
        df: Sidecar DataFrame (source_id, healpix_id, [weight])
        meta: Expected HEALPix metadata
        strict: If True, raise on violations; if False, return warnings
    
    Returns:
        dict with validation results and metrics
    
    Raises:
        AssertionError: If strict=True and validation fails
    
    Examples:
        >>> import pandas as pd
        >>> meta = HEALPyxelxMetadata(nside=32, mode='fuzzy')
        >>> sidecar = pd.DataFrame({
        ...     'source_id': [0, 0, 1],
        ...     'healpix_id': [100, 101, 100],
        ...     'weight': [0.6, 0.4, 1.0]
        ... })
        >>> result = validate_sidecar(sidecar, meta, strict=False)
        >>> result['n_unique_cells']
        2
    """
    results = {
        'valid': True,
        'errors': [],
        'warnings': [],
        'n_rows': len(df),
        'n_unique_sources': df['source_id'].nunique() if 'source_id' in df.columns else None,
        'n_unique_cells': df['healpix_id'].nunique() if 'healpix_id' in df.columns else None,
    }
    
    # Test 1: Required columns
    required_cols = ['source_id', 'healpix_id']
    missing = [c for c in required_cols if c not in df.columns]
    if missing:
        results['errors'].append(f"Missing required columns: {missing}")
        results['valid'] = False
    
    if not results['valid']:
        if strict:
            raise AssertionError(f"Sidecar validation failed: {results['errors']}")
        return results
    
    # Test 2: healpix_id range [0, npix)
    invalid_ids = (df['healpix_id'] < 0) | (df['healpix_id'] >= meta.npix)
    n_invalid = invalid_ids.sum()
    if n_invalid > 0:
        results['errors'].append(
            f"{n_invalid} healpix_id values out of range [0, {meta.npix})"
        )
        results['valid'] = False
    
    # Test 3: Unique cells ≤ npix
    if results['n_unique_cells'] > meta.npix:
        results['errors'].append(
            f"Unique cells ({results['n_unique_cells']}) exceeds npix ({meta.npix})"
        )
        results['valid'] = False
    
    # Test 4: Weight column validation (fuzzy mode)
    if meta.mode == 'fuzzy' and 'weight' in df.columns:
        # Weights must be in [0, 1]
        invalid_weights = (df['weight'] < 0) | (df['weight'] > 1)
        n_invalid_weights = invalid_weights.sum()
        if n_invalid_weights > 0:
            results['errors'].append(
                f"{n_invalid_weights} weights outside [0, 1]"
            )
            results['valid'] = False
        
        # Weights per source should sum to ~1.0 (within tolerance)
        weight_sums = df.groupby('source_id')['weight'].sum()
        tolerance = 0.01
        bad_sums = ((weight_sums < 1 - tolerance) | (weight_sums > 1 + tolerance)).sum()
        if bad_sums > 0:
            results['warnings'].append(
                f"{bad_sums} sources have weight sums != 1.0 (tolerance={tolerance})"
            )
        
        results['weight_sum_min'] = float(weight_sums.min())
        results['weight_sum_max'] = float(weight_sums.max())
    
    # Test 5: Strict mode should have 1:1 source:cell mapping
    if meta.mode == 'strict':
        multi_cell_sources = df.groupby('source_id').size()
        multi_cell = (multi_cell_sources > 1).sum()
        if multi_cell > 0:
            results['errors'].append(
                f"Strict mode: {multi_cell} sources map to multiple cells"
            )
            results['valid'] = False
    
    # Test 6: Source ID continuity (should be 0..n-1 for partition-based workflows)
    if results['n_unique_sources'] is not None:
        source_ids = df['source_id'].unique()
        expected_range = np.arange(source_ids.min(), source_ids.max() + 1)
        if not np.array_equal(np.sort(source_ids), expected_range):
            results['warnings'].append(
                "Source IDs are not contiguous (gaps detected)"
            )
    
    if strict and not results['valid']:
        raise AssertionError(f"Sidecar validation failed: {results['errors']}")
    
    return results

# %% ../nbs/00b_metadata.ipynb 8
def validate_aggregate(
    df: pd.DataFrame,
    meta: HEALPyxelxMetadata,
    expected_columns: Optional[list[str]] = None,
    strict: bool = True
) -> dict:
    """Validate aggregate DataFrame structure.
    
    Checks:
    - Index must be 'healpix_id'
    - Row count ≤ npix
    - healpix_id range [0, npix)
    - No duplicate healpix_id
    - Dense maps have sequential index [0..npix-1]
    - Expected columns present
    - n_sources column sanity
    - Finite values in stat columns
    
    Args:
        df: Aggregate DataFrame (index: healpix_id)
        meta: Expected HEALPix metadata
        expected_columns: Optional list of expected stat columns
        strict: If True, raise on violations
    
    Returns:
        dict with validation results
    
    Raises:
        AssertionError: If strict=True and validation fails
    
    Examples:
        >>> import pandas as pd
        >>> meta = HEALPyxelxMetadata(nside=32)
        >>> agg = pd.DataFrame({
        ...     'r1050_median': [0.5, 0.6],
        ...     'n_sources': [10, 15]
        ... }, index=pd.Index([0, 1], name='healpix_id'))
        >>> result = validate_aggregate(agg, meta, strict=False)
        >>> result['is_sparse']
        True
    """
    results = {
        'valid': True,
        'errors': [],
        'warnings': [],
        'n_rows': len(df),
        'is_sparse': len(df) < meta.npix,
        'is_dense': len(df) == meta.npix,
    }
    
    # Test 1: Index must be healpix_id
    if df.index.name != 'healpix_id':
        results['errors'].append(
            f"Index must be 'healpix_id', got '{df.index.name}'"
        )
        results['valid'] = False
        if strict:
            raise AssertionError(results['errors'][-1])
        return results
    
    # Test 2: Row count cannot exceed npix
    if len(df) > meta.npix:
        results['errors'].append(
            f"Row count ({len(df)}) exceeds npix ({meta.npix})"
        )
        results['valid'] = False
    
    # Test 3: healpix_id range [0, npix)
    invalid_ids = (df.index < 0) | (df.index >= meta.npix)
    n_invalid = invalid_ids.sum()
    if n_invalid > 0:
        results['errors'].append(
            f"{n_invalid} healpix_id values out of range [0, {meta.npix})"
        )
        results['valid'] = False
    
    # Test 4: No duplicate healpix_id
    if df.index.duplicated().any():
        n_dupes = df.index.duplicated().sum()
        results['errors'].append(f"{n_dupes} duplicate healpix_id values")
        results['valid'] = False
    
    # Test 5: Dense maps must have sequential index [0..npix-1]
    if results['is_dense']:
        expected_index = pd.RangeIndex(0, meta.npix, name='healpix_id')
        if not df.index.equals(expected_index):
            results['errors'].append(
                "Dense map index not aligned: expected RangeIndex(0, npix)"
            )
            results['valid'] = False
    
    # Test 6: Expected columns present
    if expected_columns:
        missing = [c for c in expected_columns if c not in df.columns]
        if missing:
            results['warnings'].append(f"Missing expected columns: {missing}")
    
    # Test 7: n_sources column sanity
    if 'n_sources' in df.columns:
        if (df['n_sources'] < 0).any():
            results['errors'].append("Negative n_sources values detected")
            results['valid'] = False
        
        results['n_sources_min'] = int(df['n_sources'].min())
        results['n_sources_max'] = int(df['n_sources'].max())
        results['n_sources_mean'] = float(df['n_sources'].mean())
    
    # Test 8: Finite value check for stat columns
    stat_cols = [c for c in df.columns if c.endswith(('_median', '_mean', '_std', '_mad'))]
    for col in stat_cols:
        n_inf = np.isinf(df[col]).sum()
        if n_inf > 0:
            results['warnings'].append(f"{n_inf} infinite values in {col}")
    
    if strict and not results['valid']:
        raise AssertionError(f"Aggregate validation failed: {results['errors']}")
    
    return results

# %% ../nbs/00b_metadata.ipynb 10
def save_with_metadata(
    df: pd.DataFrame,
    path: Union[str, Path],
    healpix_meta: HEALPyxelxMetadata,
    **kwargs
) -> None:
    """Save DataFrame with embedded HEALPix metadata.
    
    Args:
        df: DataFrame to save
        path: Output parquet path
        healpix_meta: HEALPix metadata to embed
        **kwargs: Additional arguments for pyarrow.parquet.write_table
    
    Example:
        >>> meta = HEALPyxelxMetadata(nside=32, order='nested')
        >>> save_with_metadata(agg_df, 'output.parquet', meta)
    """
    import pyarrow as pa
    
    table = pa.Table.from_pandas(df)
    
    # Merge user metadata with existing schema metadata
    existing_meta = table.schema.metadata or {}
    new_meta = {**existing_meta, **healpix_meta.to_parquet_metadata()}
    
    table = table.replace_schema_metadata(new_meta)
    pq.write_table(table, path, **kwargs)

# %% ../nbs/00b_metadata.ipynb 11
def load_with_validation(
    path: Union[str, Path],
    validate: bool = True,
    strict: bool = False
) -> tuple[pd.DataFrame, HEALPyxelxMetadata, dict]:
    """Load parquet with HEALPix metadata and optional validation.
    
    Args:
        path: Path to parquet file
        validate: If True, run validation checks
        strict: If True, raise on validation failures
    
    Returns:
        (DataFrame, HEALPyxelxMetadata, validation_results)
    
    Examples:
        >>> df, meta, results = load_with_validation('aggregate.parquet')
        >>> print(f"Valid: {results.get('valid', True)}, nside={meta.nside}")
    """
    path = Path(path)
    
    # Extract metadata
    meta = HEALPyxelxMetadata.from_parquet(path)
    
    # Load data with healpix_id as index
    df = pd.read_parquet(path)
    if 'healpix_id' in df.columns and df.index.name != 'healpix_id':
        df = df.set_index('healpix_id')
    
    # Validate if requested
    results = {}
    if validate:
        # Detect file type by structure
        if 'source_id' in df.columns or df.index.name == 'source_id':
            results = validate_sidecar(df.reset_index(), meta, strict=strict)
        else:
            results = validate_aggregate(df, meta, strict=strict)
    
    return df, meta, results

# %% ../nbs/00b_metadata.ipynb 12
def write_metadata_with_file(
    df: pd.DataFrame,
    output_path: Union[str, Path],
    file_type: FileType,
    nside: int,
    mode: str = 'fuzzy',
    order: str = 'nested',
    lon_convention: str = '0_360',
    index_name: Optional[str] = None,
    parquet_kwargs: Optional[dict] = None,
) -> HEALPyxelxMetadata:
    """Centralized metadata writer for pipeline outputs.
    
    Writes both:
    1. Embedded metadata in parquet schema (primary)
    2. Companion .meta.json file (fallback)
    
    Args:
        df: DataFrame to save
        output_path: Output parquet path
        file_type: FileType enum indicating pipeline stage
        nside: HEALPix resolution parameter
        mode: Cell assignment mode ('strict' or 'fuzzy')
        order: Pixel ordering ('nested' or 'ring')
        lon_convention: Longitude convention ('0_360' or '-180_180')
        index_name: Optional name for DataFrame index (e.g., 'healpix_id')
        parquet_kwargs: Additional arguments for pyarrow.parquet.write_table
    
    Returns:
        HEALPyxelxMetadata instance written to file
    
    Examples:
        >>> meta = write_metadata_with_file(
        ...     agg_df, 
        ...     'output.parquet',
        ...     FileType.AGGREGATE,
        ...     nside=32,
        ...     mode='fuzzy',
        ...     order='nested'
        ... )
    """
    import pyarrow as pa
    
    output_path = Path(output_path)
    parquet_kwargs = parquet_kwargs or {}
    
    # Create metadata object
    metadata = HEALPyxelxMetadata(
        nside=nside,
        order=order,
        mode=mode,
        lon_convention=lon_convention,
        file_type=file_type,
    )
    
    # Convert DataFrame to Arrow table
    table = pa.Table.from_pandas(df)
    
    # Embed metadata in schema
    existing_meta = table.schema.metadata or {}
    new_meta = {**existing_meta, **metadata.to_parquet_metadata()}
    table = table.replace_schema_metadata(new_meta)
    
    # Write parquet with embedded metadata
    pq.write_table(table, output_path, **parquet_kwargs)
    
    # Also write companion .meta.json for backwards compatibility
    json_path = output_path.with_suffix('.meta.json')
    with open(json_path, 'w') as f:
        json.dump(metadata.to_dict(), f, indent=2)
    
    return metadata

# %% ../nbs/00b_metadata.ipynb 13
def read_metadata_from_file(
    path: Union[str, Path],
    fallback_to_json: bool = True,
) -> tuple[HEALPyxelxMetadata, Optional[FileType]]:
    """Unified metadata reader with intelligent fallback strategy.
    
    Tries parquet schema metadata first (embedded by write_metadata_with_file),
    then falls back to .meta.json if present.
    
    Args:
        path: Path to parquet file
        fallback_to_json: If True, fall back to .meta.json; if False, raise
    
    Returns:
        (HEALPyxelxMetadata, FileType) tuple
    
    Raises:
        ValueError: If no metadata found and fallback_to_json=False
    
    Examples:
        >>> meta, file_type = read_metadata_from_file('aggregate.parquet')
        >>> print(f"Type: {file_type}, nside={meta.nside}")
    """
    path = Path(path)
    
    # Try embedded parquet metadata first
    try:
        pf = pq.ParquetFile(path)
        schema_meta = pf.schema_arrow.metadata
        if schema_meta and b'healpix_metadata' in schema_meta:
            meta_json = schema_meta[b'healpix_metadata'].decode()
            meta = HEALPyxelxMetadata.from_dict(json.loads(meta_json))
            return meta, meta.file_type
    except Exception as e:
        pass
    
    # Fallback to companion .meta.json
    json_path = path.with_suffix('.meta.json')
    if json_path.exists():
        with open(json_path) as f:
            meta_dict = json.load(f)
            meta = HEALPyxelxMetadata.from_dict(meta_dict)
            return meta, meta.file_type
    
    if not fallback_to_json:
        raise ValueError(
            f"No HEALPix metadata found for {path} "
            f"(no parquet schema metadata, no {json_path.name})"
        )
    
    raise ValueError(f"No HEALPix metadata found for {path}")

# %% ../nbs/00b_metadata.ipynb 14
def validate_file_type(
    path: Union[str, Path],
    expected_type: FileType,
    strict: bool = False,
) -> dict:
    """Validate that a file matches expected pipeline stage.
    
    Checks:
    - File exists
    - Metadata present and parseable
    - file_type field matches expected value
    
    Args:
        path: Path to parquet file
        expected_type: Expected FileType
        strict: If True, raise on mismatch; if False, return results dict
    
    Returns:
        dict with validation results
    
    Raises:
        ValueError: If strict=True and validation fails
    
    Examples:
        >>> results = validate_file_type('agg.parquet', FileType.AGGREGATE)
        >>> print(f"Valid: {results['valid']}, found_type={results['found_type']}")
    """
    path = Path(path)
    results = {
        'valid': False,
        'path': str(path),
        'expected_type': expected_type.value,
        'found_type': None,
        'errors': [],
        'warnings': [],
    }
    
    # Test 1: File exists
    if not path.exists():
        results['errors'].append(f"File does not exist: {path}")
        if strict:
            raise ValueError(results['errors'][-1])
        return results
    
    # Test 2: Read metadata
    try:
        meta, found_type = read_metadata_from_file(path)
        results['found_type'] = found_type.value if found_type else None
    except Exception as e:
        results['errors'].append(f"Could not read metadata: {e}")
        if strict:
            raise ValueError(results['errors'][-1])
        return results
    
    # Test 3: Match expected type
    if found_type != expected_type:
        results['errors'].append(
            f"File type mismatch: expected {expected_type.value}, "
            f"found {found_type.value if found_type else 'None'}"
        )
        if strict:
            raise ValueError(results['errors'][-1])
        return results
    
    results['valid'] = True
    return results

# %% ../nbs/00b_metadata.ipynb 24
def load_cli_aggregate(
    cli_output_dir: Union[str, Path],
    filename_pattern: str = '*sparse_aggregate.parquet',
    validate: bool = True,
    strict: bool = False
) -> tuple[pd.DataFrame, HEALPyxelxMetadata, dict]:
    """Load and validate CLI-generated aggregate from standard output directory.
    
    Args:
        cli_output_dir: Path to CLI output directory
        filename_pattern: Glob pattern for aggregate files
        validate: If True, run validation
        strict: If True, raise on validation failures
    
    Returns:
        (DataFrame, HEALPyxelxMetadata, validation_results)
    
    Raises:
        FileNotFoundError: If no matching files found
        ValueError: If multiple ambiguous matches
    """
    cli_output_dir = Path(cli_output_dir)
    files = list(cli_output_dir.glob(filename_pattern))
    if not files:
        raise FileNotFoundError(f"No files matching '{filename_pattern}' in {cli_output_dir}")
    if len(files) > 1:
        raise ValueError(f"Ambiguous match: {len(files)} files match '{filename_pattern}'")
    return load_with_validation(files[0], validate=validate, strict=strict)
