Sentinel-2¶

# Optimal bands for aquatic applications
WATER_BANDS = {
    'blue': 'B2',        # Deep penetration, bathymetry
    'green': 'B3',       # Balance of penetration/sensitivity
    'red': 'B4',         # Shallow waters, SPM
    'red_edge': 'B5',    # Chlorophyll, advanced algorithms
    'nir': 'B8',         # High turbidity, SPM
    'nir_narrow': 'B8A', # Water vapour, atmospheric correction
    'swir1': 'B11',      # Cloud/snow detection
}

# NDVI (Normalized Difference Vegetation Index)
def calculate_ndvi(image):
    """NDVI = (NIR - Red) / (NIR + Red)"""
    nir = image.select('B8')
    red = image.select('B4')
    ndvi = nir.subtract(red).divide(nir.add(red)).rename('NDVI')
    return image.addBands(ndvi)

# EVI (Enhanced Vegetation Index)
def calculate_evi(image):
    """EVI = 2.5 * (NIR - Red) / (NIR + 6*Red - 7.5*Blue + 1)"""
    nir = image.select('B8')
    red = image.select('B4')
    blue = image.select('B2')

    evi = nir.subtract(red).multiply(2.5).divide(
        nir.add(red.multiply(6)).subtract(blue.multiply(7.5)).add(1)
    ).rename('EVI')
    return image.addBands(evi)

# NDWI (Normalized Difference Water Index)
def calculate_ndwi(image):
    """NDWI = (Green - NIR) / (Green + NIR)"""
    green = image.select('B3')
    nir = image.select('B8')
    ndwi = green.subtract(nir).divide(green.add(nir)).rename('NDWI')
    return image.addBands(ndwi)

# MNDWI (Modified NDWI)
def calculate_mndwi(image):
    """MNDWI = (Green - SWIR1) / (Green + SWIR1)"""
    green = image.select('B3')
    swir = image.select('B11')
    mndwi = green.subtract(swir).divide(green.add(swir)).rename('MNDWI')
    return image.addBands(mndwi)

from gee_acolite.sensors.sentinel2 import BAND_BY_SCALE

def get_native_projection(scale):
    """
    Get the representative band for a given scale.

    Parameters
    ----------
    scale : int
        Target scale (10, 20, or 60 metres)

    Returns
    -------
    str
        Band name with that native resolution
    """
    return BAND_BY_SCALE.get(scale, 'B2')

# Example: resample all bands to 10m
target_scale = 10
reference_band = get_native_projection(target_scale)

def resample_to_10m(image):
    """Resample all bands to 10m."""
    proj = image.select(reference_band).projection()
    return image.resample('bilinear').reproject(crs=proj, scale=target_scale)

# S2A/S2B harmonisation (simplified method)
def harmonize_s2(image):
    """
    Apply offset correction between S2A and S2B.
    Based on Roy et al. (2021).
    """
    # Identify satellite
    spacecraft = image.get('SPACECRAFT_NAME')

    # Offsets for S2B (S2A is the reference)
    offsets_s2b = {
        'B2': 0.0008,   # Blue
        'B3': -0.0001,  # Green
        'B4': -0.0017,  # Red
        'B8': -0.0022,  # NIR
        'B11': -0.0011, # SWIR1
        'B12': -0.0028  # SWIR2
    }

    # Apply offset if S2B
    def apply_offset(band_name):
        band = image.select(band_name)
        offset = offsets_s2b.get(band_name, 0)
        return ee.Image(
            ee.Algorithms.If(
                ee.String(spacecraft).equals('Sentinel-2B'),
                band.subtract(offset),
                band
            )
        ).rename(band_name)

    # Apply to all bands
    harmonized_bands = [apply_offset(b) for b in offsets_s2b.keys()]
    other_bands = image.select([b for b in image.bandNames().getInfo()
                                if b not in offsets_s2b.keys()])

    return ee.Image.cat(harmonized_bands + [other_bands]).copyProperties(image)

# Common RGB composites

# True Colour
def true_color(image):
    """RGB: B4-B3-B2 (Red-Green-Blue)"""
    return image.select(['B4', 'B3', 'B2'])

# False Colour (NIR)
def false_color_nir(image):
    """RGB: B8-B4-B3 (NIR-Red-Green)"""
    return image.select(['B8', 'B4', 'B3'])

# False Colour (SWIR)
def false_color_swir(image):
    """RGB: B12-B8A-B4 (SWIR2-NIR-Red)"""
    return image.select(['B12', 'B8A', 'B4'])

# Agriculture
def agriculture(image):
    """RGB: B11-B8-B2 (SWIR1-NIR-Blue)"""
    return image.select(['B11', 'B8', 'B2'])

# Bathymetry
def bathymetry_rgb(image):
    """RGB: B4-B3-B2 with aggressive stretch for water."""
    rgb = image.select(['B4', 'B3', 'B2'])
    return rgb.multiply(3).clamp(0, 1)