src/data/netcdf_loader.py

"""
NetCDF Data Loader

Memory-efficient loader for ERA5 and synthetic climate data in NetCDF format.
Uses chunked loading and lazy evaluation to handle datasets larger than RAM.

Design Decisions:
- xarray with dask backend for lazy evaluation
- Iterator pattern for memory-bounded processing
- Automatic coordinate normalization
- Support for both real ERA5 and synthetic data

Why Chunked Loading?
- ERA5 global data at 0.25° can exceed 100GB
- India region for 1 year ≈ 3-5GB
- 8GB RAM requires streaming, not full load

Time Complexity: O(n) for n timesteps, each loaded once
Space Complexity: O(chunk_size) - constant memory usage
"""

import numpy as np
from pathlib import Path
from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
import warnings


class NetCDFLoader:
    """
    Memory-efficient NetCDF data loader for climate data.
    
    Supports ERA5 reanalysis data and synthetic test data.
    Provides chunked iteration for large datasets.
    
    Attributes:
        filepath: Path to NetCDF file
        config: Configuration dictionary
        dataset: xarray Dataset (lazy loaded)
    """
    
    # ERA5 variable name mappings (ERA5 name -> standard name)
    ERA5_VARIABLE_MAP = {
        "t2m": "t2m",           # 2m temperature
        "2t": "t2m",            # Alternative ERA5 name
        "tp": "tp",             # Total precipitation
        "r": "r",               # Relative humidity
        "rh": "r",              # Alternative name
    }
    
    def __init__(
        self,
        filepath: Union[str, Path],
        config: Dict[str, Any],
        lazy: bool = True
    ):
        """
        Initialize the NetCDF loader.
        
        Args:
            filepath: Path to NetCDF file
            config: Configuration dictionary
            lazy: If True, use lazy loading (recommended for large files)
        """
        self.filepath = Path(filepath)
        self.config = config
        self.lazy = lazy
        self.dataset = None
        self._is_loaded = False
        
        # Extract region config
        region = config.get("data", {}).get("region", {})
        self.lat_min = region.get("lat_min", 6.0)
        self.lat_max = region.get("lat_max", 38.0)
        self.lon_min = region.get("lon_min", 68.0)
        self.lon_max = region.get("lon_max", 98.0)
        
        # Variables to load
        self.variables = config.get("data", {}).get("variables", ["t2m"])
    
    def load(self) -> None:
        """
        Load the NetCDF dataset.
        
        Uses lazy loading by default to avoid loading entire file into memory.
        """
        try:
            import xarray as xr
        except ImportError:
            raise ImportError("xarray required: pip install xarray netCDF4")
        
        if not self.filepath.exists():
            raise FileNotFoundError(f"NetCDF file not found: {self.filepath}")
        
        # Standard loading without dask chunking for simplicity
        # This works without dask installed
        self.dataset = xr.open_dataset(self.filepath)
        
        # Normalize coordinate names
        self._normalize_coordinates()
        
        # Normalize variable names
        self._normalize_variables()
        
        # Subset to region of interest
        self._subset_region()
        
        self._is_loaded = True
    
    def _normalize_coordinates(self) -> None:
        """
        Normalize coordinate names to standard format.
        
        ERA5 files may use 'lat'/'lon' or 'latitude'/'longitude'.
        This method ensures consistent naming.
        """
        rename_map = {}
        
        # Latitude normalization
        for name in ["lat", "latitude", "Latitude", "LAT"]:
            if name in self.dataset.coords and name != "latitude":
                rename_map[name] = "latitude"
        
        # Longitude normalization
        for name in ["lon", "longitude", "Longitude", "LON"]:
            if name in self.dataset.coords and name != "longitude":
                rename_map[name] = "longitude"
        
        if rename_map:
            self.dataset = self.dataset.rename(rename_map)

    def _normalize_variables(self) -> None:
        """
        Normalize variable names to internal standard.
        
        Maps various ERA5 naming conventions (e.g. '2t', 'var167')
        to standard project names (e.g. 't2m').
        """
        rename_map = {}
        
        # Temperature mapping
        for name in ["2t", "var167", "2m_temperature", "VAR_2T"]:
            if name in self.dataset.data_vars:
                rename_map[name] = "t2m"
                
        # Precipitation mapping
        for name in ["tp", "var228", "total_precipitation", "VAR_TP"]:
            if name in self.dataset.data_vars and name != "tp":
                rename_map[name] = "tp"

        if rename_map:
            self.dataset = self.dataset.rename(rename_map)
    
    def _subset_region(self) -> None:
        """
        Subset dataset to the region of interest.
        
        This is where most memory savings occur - we only keep
        the India bounding box, discarding ~95% of global data.
        """
        # Get current lat/lon ranges
        lats = self.dataset.coords["latitude"].values
        lons = self.dataset.coords["longitude"].values
        
        # Handle longitude wrapping (0-360 vs -180-180)
        if lons.max() > 180:
            # Convert 0-360 to -180-180 if needed
            lon_min = self.lon_min if self.lon_min >= 0 else self.lon_min + 360
            lon_max = self.lon_max if self.lon_max >= 0 else self.lon_max + 360
        else:
            lon_min, lon_max = self.lon_min, self.lon_max
        
        # Determine lat order (some files are N-to-S)
        lat_ascending = lats[0] < lats[-1]
        
        if lat_ascending:
            lat_slice = slice(self.lat_min, self.lat_max)
        else:
            lat_slice = slice(self.lat_max, self.lat_min)
        
        # Apply subsetting
        self.dataset = self.dataset.sel(
            latitude=lat_slice,
            longitude=slice(lon_min, lon_max)
        )
    
    def get_variable(self, var_name: str) -> np.ndarray:
        """
        Get a complete variable as a NumPy array.
        
        Warning: This loads the entire variable into memory.
        Use iterate_chunks() for large datasets.
        
        Args:
            var_name: Variable name (e.g., 't2m', 'tp')
            
        Returns:
            Array of shape (time, lat, lon)
        """
        if not self._is_loaded:
            self.load()
        
        # Map to standard name if needed
        mapped_name = self.ERA5_VARIABLE_MAP.get(var_name, var_name)
        
        if mapped_name not in self.dataset.data_vars:
            available = list(self.dataset.data_vars)
            raise KeyError(
                f"Variable '{var_name}' not found. Available: {available}"
            )
        
        data = self.dataset[mapped_name].values
        
        # Ensure float32 for memory efficiency
        return data.astype(np.float32)
    
    def iterate_chunks(
        self,
        var_name: str,
        chunk_size: int = 10
    ) -> Iterator[Tuple[np.ndarray, int]]:
        """
        Iterate over temporal chunks of a variable.
        
        Memory-efficient iteration for large datasets.
        Each chunk contains chunk_size consecutive timesteps.
        
        Args:
            var_name: Variable name
            chunk_size: Number of timesteps per chunk
            
        Yields:
            Tuple of (data_chunk, start_index)
        """
        if not self._is_loaded:
            self.load()
        
        mapped_name = self.ERA5_VARIABLE_MAP.get(var_name, var_name)
        var_data = self.dataset[mapped_name]
        
        n_times = len(self.dataset.coords["time"])
        
        for start_idx in range(0, n_times, chunk_size):
            end_idx = min(start_idx + chunk_size, n_times)
            
            # Select time slice
            chunk = var_data.isel(time=slice(start_idx, end_idx))
            
            # Load into memory (compute if dask-backed)
            chunk_data = chunk.values.astype(np.float32)
            
            yield chunk_data, start_idx
    
    def get_coordinates(self) -> Tuple[np.ndarray, np.ndarray]:
        """
        Get latitude and longitude coordinate arrays.
        
        Returns:
            Tuple of (latitudes, longitudes) arrays
        """
        if not self._is_loaded:
            self.load()
        
        lats = self.dataset.coords["latitude"].values.astype(np.float32)
        lons = self.dataset.coords["longitude"].values.astype(np.float32)
        
        return lats, lons
    
    def get_time_range(self) -> Tuple[Any, Any]:
        """
        Get the temporal range of the dataset.
        
        Returns:
            Tuple of (start_time, end_time)
        """
        if not self._is_loaded:
            self.load()
        
        times = self.dataset.coords["time"].values
        return times[0], times[-1]
    
    def get_shape(self) -> Dict[str, int]:
        """
        Get dataset dimensions.
        
        Returns:
            Dictionary with 'time', 'latitude', 'longitude' sizes
        """
        if not self._is_loaded:
            self.load()
        
        return {
            "time": len(self.dataset.coords["time"]),
            "latitude": len(self.dataset.coords["latitude"]),
            "longitude": len(self.dataset.coords["longitude"]),
        }
    
    def get_variable_stats(self, var_name: str) -> Dict[str, float]:
        """
        Compute statistics for a variable (streaming computation).
        
        Uses Welford's online algorithm for memory-efficient stats.
        
        Args:
            var_name: Variable name
            
        Returns:
            Dictionary with 'mean', 'std', 'min', 'max'
        """
        # Online statistics computation
        n = 0
        mean = 0.0
        M2 = 0.0
        min_val = float("inf")
        max_val = float("-inf")
        
        for chunk, _ in self.iterate_chunks(var_name):
            for x in chunk.flatten():
                n += 1
                delta = x - mean
                mean += delta / n
                delta2 = x - mean
                M2 += delta * delta2
                min_val = min(min_val, x)
                max_val = max(max_val, x)
        
        variance = M2 / n if n > 1 else 0.0
        std = np.sqrt(variance)
        
        return {
            "mean": float(mean),
            "std": float(std),
            "min": float(min_val),
            "max": float(max_val),
            "n_samples": n,
        }
    
    def close(self) -> None:
        """Close the dataset and free resources."""
        if self.dataset is not None:
            self.dataset.close()
            self.dataset = None
            self._is_loaded = False
    
    def __enter__(self):
        """Context manager entry."""
        self.load()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.close()
        return False


def load_climate_data(
    config: Dict[str, Any],
    data_path: Optional[str] = None
) -> NetCDFLoader:
    """
    Convenience function to load climate data based on configuration.
    
    Automatically detects data source (synthetic vs local) from config.
    
    Args:
        config: Configuration dictionary
        data_path: Optional explicit path to data file
        
    Returns:
        Initialized NetCDFLoader
    """
    if data_path is not None:
        filepath = Path(data_path)
    else:
        source = config.get("data", {}).get("source", "synthetic")
        raw_dir = config.get("data", {}).get("raw_dir", "data/raw")
        
        if source == "synthetic":
            filepath = Path(raw_dir) / "synthetic_climate.nc"
        elif source == "cds_api":
            # Look for ERA5 files specifically
            filepath = Path(raw_dir)
            nc_files = list(filepath.glob("era5*.nc"))
            if not nc_files:
                raise FileNotFoundError(
                    f"No ERA5 files (era5*.nc) found in {raw_dir}. "
                    "Run python src/data/cds_fetcher.py first."
                )
            # Pick the most recent one (presumably what we want)
            filepath = sorted(nc_files)[-1]
        else:
            # Look for any NetCDF files
            filepath = Path(raw_dir)
            nc_files = list(filepath.glob("*.nc"))
            if not nc_files:
                raise FileNotFoundError(
                    f"No NetCDF files found in {raw_dir}. "
                    "Run generate_synthetic_data.py first."
                )
            filepath = nc_files[0]
    
    loader = NetCDFLoader(filepath, config)
    loader.load()
    
    return loader