src/models/land_mask.py

"""
Land-Sea Mask for ERA5 India domain.

IMPROVEMENT 4: Land-sea mask
The AI model (Swin2SR) has no knowledge of land vs ocean.
Applying super-resolution to ocean pixels creates hallucinated
fine-scale SST structure that doesn't exist physically.

This module builds a land mask at both ERA5 resolution (129x121)
and output resolution (516x484) using two methods:

Method 1 - Temperature variance proxy:
  Ocean pixels have very low temporal variance (SST is smooth and slow).
  Land pixels have high variance (diurnal cycle, seasons, weather).
  Threshold the annual std dev to separate land from sea.

Method 2 - Known ocean bounds (backup):
  Hard-code the approximate lat/lon rectangles of:
  - Arabian Sea (west coast)
  - Bay of Bengal (east coast)
  - Indian Ocean (south)
  Use as fallback or cross-validation.

The mask is applied after AI inference:
  final_output[ocean] = ERA5_upsampled[ocean]   (original SST preserved)
  final_output[land]  = AI_prediction[land]      (AI enhancement applied)

This is scientifically correct: ERA5 SST is already high quality
at 0.25deg. Only land surface temperatures benefit from SR enhancement.
"""
import numpy as np
from scipy.ndimage import zoom as spz, binary_dilation


def build_land_mask(data: np.ndarray,
                    lat_min: float, lat_max: float,
                    lon_min: float, lon_max: float,
                    output_scale: int = 4) -> tuple:
    """
    Build land mask from temporal variance of ERA5 data.

    Returns:
        mask_lr  : (H, W)   bool  True=land at ERA5 resolution
        mask_hr  : (4H, 4W) bool  True=land at output resolution
    """
    T, H, W = data.shape

    # Temporal std dev at each grid point across all 8784 timesteps
    std_map = data.std(axis=0)   # (H, W)

    # Ocean = low variance (SST changes slowly)
    # Land  = high variance (weather, diurnal, seasonal)
    # Use Otsu-like threshold: midpoint between min and median
    thresh = (std_map.min() + np.percentile(std_map, 40)) / 2
    land_mask_lr = std_map > thresh   # True = land

    # Morphological cleanup — remove isolated ocean pixels on coast
    land_mask_lr = binary_dilation(land_mask_lr, iterations=1)

    # Hard-mask: known open ocean bounding boxes
    lats = np.linspace(lat_max, lat_min, H)   # row 0 = north
    lons = np.linspace(lon_min, lon_max, W)

    ocean_boxes = [
        # (lat_north, lat_south, lon_west, lon_east)
        # Conservative boxes — only clear open ocean, not coastal zones
        # Coastal land (Gujarat, Kerala, AP coast) handled by variance method
        (38.0,  6.0, 68.0, 71.0),   # Arabian Sea open water (far west)
        (12.0,  6.0, 71.0, 79.0),   # Indian Ocean (south tip, below tip of India)
        (22.0,  6.0, 88.0, 98.0),   # Bay of Bengal (far east open water)
    ]
    for ln, ls, lw, le in ocean_boxes:
        for r in range(H):
            for c in range(W):
                if ls <= lats[r] <= ln and lw <= lons[c] <= le:
                    land_mask_lr[r, c] = False   # force ocean

    # Upsample to output resolution
    land_mask_hr_f = spz(land_mask_lr.astype(float), output_scale, order=0)
    land_mask_hr = land_mask_hr_f > 0.5

    n_land = land_mask_lr.sum()
    n_total = H * W
    print(f"  [Mask] Land pixels: {n_land}/{n_total} "
          f"({n_land/n_total*100:.1f}%) at ERA5 resolution")
    print(f"  [Mask] Output mask: {land_mask_hr.shape}")

    return land_mask_lr, land_mask_hr


def apply_land_mask(pred_K: np.ndarray,
                    era5_K: np.ndarray,
                    mask_hr: np.ndarray) -> np.ndarray:
    """
    Blend AI prediction with ERA5 using land-sea mask.

    pred_K  : (4H, 4W) AI prediction in Kelvin
    era5_K  : (H, W)   ERA5 input in Kelvin
    mask_hr : (4H, 4W) True=land, False=ocean

    Returns: (4H, 4W) masked output
    """
    # Upsample ERA5 to output resolution for ocean pixels
    H, W = era5_K.shape
    era5_up = spz(era5_K, pred_K.shape[0] / H, order=3)
    era5_up = era5_up[:pred_K.shape[0], :pred_K.shape[1]]

    # Blend: land=AI, ocean=ERA5
    out = era5_up.copy()
    out[mask_hr] = pred_K[mask_hr]
    return out