Data Cleaning & Preprocessing

# Import essential libraries

import h5py
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

# ---- setup -------------------------------------------------------------

def download_tanager_data(url, file_path):

    import os
    import urllib.request

    def _progress(count, block_size, total_size):
        mb_done = count * block_size / 1_000_000
        mb_total = total_size / 1_000_000
        print(f"\rDownloading... {mb_done:.1f} / {mb_total:.1f} MB", end="", flush=True)

    if not os.path.exists(file_path):
        urllib.request.urlretrieve(url, file_path, reporthook=_progress)
        print(f"\nDownload complete: {file_path}")
    else:
        print(f"File already exists, skipping download: {file_path}")
    return file_path

# ---- io ----------------------------------------------------------------

def print_structure(name, obj):

    level = name.count('/')
    indent = '    ' * level

    if isinstance(obj, h5py.Group):
        print(f"{indent}📂 {name}/")
    elif isinstance(obj, h5py.Dataset):
        print(f"{indent}📄 {name} | {obj.shape} | {obj.dtype}")

    if len(obj.attrs) > 0:
        for key, val in obj.attrs.items():
            val_str = str(val)
            if len(val_str) > 50:
                val_str = val_str[:50] + "..."
            print(f"{indent}    ↳ 🏷️ {key}: {val_str}")

def load_hdf5(file_path,
    data_path='HDFEOS/SWATHS/HYP/Data Fields/surface_reflectance',
    band=None
    ):

    with h5py.File(file_path, 'r') as f:
        dset = f[data_path]
        data = dset[:] if band is None else dset[band]
        attrs = dict(dset.attrs)
    return data, attrs

# ---- bands -------------------------------------------------------------

def nearest_band(wavelengths, target_nm):
    """Return the index of the band whose wavelength is closest to ``target_nm``.

    Parameters
    ----------
    wavelengths : ndarray
        Per-band center wavelengths (nm).
    target_nm : float
        Wavelength of interest (nm).

    Returns
    -------
    int
        Index of the nearest band.
    """
    return int(np.argmin(np.abs(np.asarray(wavelengths) - target_nm)))

# ---- visualization ----------------------------------------------------

def make_rgb(cube, wavelengths, r_nm=680, g_nm=560, b_nm=470,
             percentile_low=2, percentile_high=98):
    """Build a contrast-stretched true-colour RGB composite from a cube.

    Picks the bands nearest the three target wavelengths, stacks them, and
    applies a percentile stretch (computed with ``nanpercentile`` so NaN
    no-data is ignored).

    Parameters
    ----------
    cube : ndarray, shape (bands, rows, cols)
        Hyperspectral cube.
    wavelengths : ndarray
        Per-band center wavelengths (nm).
    r_nm, g_nm, b_nm : float
        Target wavelengths (nm) for the red, green, blue channels.
    percentile_low, percentile_high : float
        Percentiles for the contrast stretch.

    Returns
    -------
    ndarray, shape (rows, cols, 3)
        RGB image normalized to [0, 1].
    """
    r = nearest_band(wavelengths, r_nm)
    g = nearest_band(wavelengths, g_nm)
    b = nearest_band(wavelengths, b_nm)
    rgb = np.dstack((cube[r], cube[g], cube[b]))
    p_low, p_high = np.nanpercentile(rgb, (percentile_low, percentile_high))
    return np.clip((rgb - p_low) / (p_high - p_low), 0, 1)

# Download the Tanager-1 Ortho SR HDF5 file (a cloudy scene) using your Toolkit's
# download_tanager_data (defined in the Toolkit cell above).
url = "https://storage.googleapis.com/open-cogs/planet-stac/tanager1-release2-core-imagery/ortho_sr_hdf5/20250301_143913_32_4001_ortho_sr_hdf5.h5"

file_path = '20250301_143913_32_4001_ortho_sr_hdf5.h5'

download_tanager_data(url, file_path)

# HDF5-EOS Internal Paths (Tanager-1 Standard)
path_sr  = "HDFEOS/GRIDS/HYP/Data Fields/surface_reflectance"

# Quality Masks
path_cloud  = "HDFEOS/GRIDS/HYP/Data Fields/beta_cloud_mask"
path_cirrus = "HDFEOS/GRIDS/HYP/Data Fields/beta_cirrus_mask"
path_nodata = "HDFEOS/GRIDS/HYP/Data Fields/nodata_pixels"

# 1. Load the cube + wavelengths with your Toolkit loader (simple version, raw cube)
sr_cube, attrs = load_hdf5(file_path, path_sr)

good_bands = attrs['good_wavelengths']
wavelengths = attrs['wavelengths']

# 2. Read the QA layers this lesson needs (the simple loader doesn't return them yet)
mask_cloud, attrs = load_hdf5(file_path, path_cloud)
mask_cirrus, attrs = load_hdf5(file_path, path_cirrus)
mask_nodata, attrs = load_hdf5(file_path, path_nodata)

# 3. Convert Fill Values to NaN immediately to prevent plotting errors
fill_value = attrs['_FillValue']
sr_cube[sr_cube == fill_value] = np.nan

# 1. Check Data Dimensions
# For Tanager-1, the shape is usually (Bands, Lines, Samples)
# though some APIs might transpose it.
n_bands, n_lines, n_samples = sr_cube.shape

print(f"--- Data Dimensions ---")
print(f"Number of Bands: {n_bands}")
print(f"Spatial Resolution (Lines x Samples): {n_lines} x {n_samples}")
print(f"Total Pixels: {n_lines * n_samples:,}")

# 2. Inspect Spectral Range
min_wl = np.min(wavelengths)
max_wl = np.max(wavelengths)
mean_spacing = np.mean(np.diff(wavelengths))

print(f"\n--- Spectral Information ---")
print(f"Wavelength Range: {min_wl:.2f} nm to {max_wl:.2f} nm")
print(f"Average Spectral Sampling: {mean_spacing:.2f} nm")

# 3. Identify 'Good' Bands
# Tanager uses a boolean or binary mask for bands unaffected by water vapor/noise
num_good = np.sum(good_bands)
print(f"Number of 'Good' Bands (Non-Atmospheric Absorption): {num_good} / {n_bands}")

# 4. Check Mask Statistics
# Understanding how much of your tile is actually usable (not cloud or nodata)
def get_mask_pct(mask):
    return (np.sum(mask > 0) / mask.size) * 100

print(f"\n--- Quality Assessment ---")
print(f"Cloud Cover: {get_mask_pct(mask_cloud):.2f}%")
print(f"Cirrus Cover: {get_mask_pct(mask_cirrus):.2f}%")
print(f"NoData Pixels: {get_mask_pct(mask_nodata):.2f}%")

print(f"Original Cube Shape: {sr_cube.shape} (Bands, Height, Width)")

# 1. Identify valid indices
# We look for indices where the 'good_wavelengths' flag is set to 1.
valid_band_indices = np.where(good_bands == 1)[0]

# 2. Slice the Cube and Wavelengths
# We subset the cube along the first dimension (Bands)
# The spatial dimensions (Height, Width) remain untouched.
sr_cube_clean = sr_cube[valid_band_indices, :, :]
wavelengths_clean = wavelengths[valid_band_indices]

# 3. Verify the result
bands_removed = sr_cube.shape[0] - sr_cube_clean.shape[0]

print(f"Bands Removed: {bands_removed}")
print(f"New Cube Shape:  {sr_cube_clean.shape}")
print(f"Remaining Wavelengths: {len(wavelengths_clean)}")

# 1.Create the Composite Mask
# We use the Bitwise OR operator (|) which functions as the Union (∪)
combined_mask = (mask_cloud == 1) | (mask_cirrus == 1) | (mask_nodata == 1)

# 2. Apply the Mask
# We set invalid pixels to NaN (Not a Number).
# Note on Broadcasting:
# The mask is 2D (Lines, Samples). The cube is 3D (Bands, Lines, Samples).
# NumPy automatically applies this 2D mask to EVERY band in the cube.
sr_cube_masked = sr_cube_clean.copy()
sr_cube_masked[:, combined_mask] = np.nan

# 3. Calculate the percentage of true masked values
total_cloud_pixels = np.sum((mask_cloud == 1) | (mask_cirrus == 1))
total_pixels = sr_cube_masked.shape[1] * sr_cube_masked.shape[2]
true_total_pixels = total_pixels - np.sum(mask_nodata)
true_cloud_pct = (total_cloud_pixels / true_total_pixels) * 100
print(f"True Cloud Cover: {true_cloud_pct:.2f}%")

# Generate RGB composite and paint masked pixels yellow for visualization
rgb_masked = make_rgb(sr_cube_masked, wavelengths_clean, r_nm=680, g_nm=550, b_nm=470)
rgb_masked[combined_mask] = [1, 1, 0]  # Mark masked pixels as yellow

fig, axes = plt.subplots(1, 2, figsize=(16, 8), constrained_layout=True)

# Left: Show combined mask
im_mask = axes[0].imshow(combined_mask)
axes[0].set_title("Step 1: The Combined Mask\n(Yellow = Pixels to Remove)", fontsize=14)
axes[0].axis('off')

# Legend for mask plot
patches = [mpatches.Patch(color=im_mask.cmap(im_mask.norm(val)), label=lab) for val, lab in zip([0,1], ["Valid Data", "Masked Area"])]
axes[0].legend(handles=patches, loc='lower right', borderaxespad=0.5)

# Right: Show masked RGB composite
axes[1].imshow(rgb_masked)
axes[1].set_title("Step 2: Final RGB Composite\n(Yellow areas indicate removed data)", fontsize=14)
axes[1].axis('off')

plt.suptitle("Masking Validation: Visualizing Data Removal", fontsize=16, weight='bold')
plt.show()

# =====================================================================
# 🧰 TOOLKIT — YOUR TURN  (implement before the next lesson)
# Fill in the body below; the reference solution appears at the top
# of the next coding lesson (Module 4 · Lesson 1, Part 2: Spectral Indices).
# =====================================================================
import h5py
import numpy as np

# ---- io ----------------------------------------------------------------

def load_hdf5(
    file_path,
    data_path='HDFEOS/GRIDS/HYP/Data Fields/surface_reflectance',
    nodata_path='HDFEOS/GRIDS/HYP/Data Fields/nodata_pixels',
    clean=True,
    band=None
    ):
    """Load a Tanager-1 hyperspectral cube (cleaning version).

    Upgraded in Module 3 · Lesson 1. With ``clean=True`` the two universal
    cleaning steps are applied: bad bands (``good_wavelengths == 0``) are
    dropped and ``nodata_pixels`` are set to NaN. Pass ``band=`` to read a
    single 2-D band straight from disk instead of the whole cube.

    Parameters
    ----------
    file_path : str
        Path to the Tanager-1 ``.h5`` file.
    data_path : str, optional
        Internal HDF5 path to the reflectance dataset (default: Ortho / GRIDS).
    nodata_path : str, optional
        Internal HDF5 path to the ``nodata_pixels`` mask.
    clean : bool, optional
        If True, drop bad bands and mask no-data pixels to NaN.
    band : int, optional
        If given, read only that band (a single 2-D slice) instead of the
        full cube — a quick, memory-light look at one wavelength.

    Returns
    -------
    data : ndarray
        Reflectance values: the full cube ``(bands, rows, cols)`` when
        ``band`` is None, otherwise a single ``(rows, cols)`` band. Bad
        bands are removed and no-data set to NaN when ``clean=True``.
    attrs : dict
        Dataset attributes, with ``'wavelengths'`` and ``'good_wavelengths'``
        trimmed to match ``data``.
    no_data : ndarray, shape (rows, cols), dtype=bool
        Boolean no-data mask.
    """
    # ------------------------------------------------------------------
    # HINTS — replace `pass` with your implementation
    # ------------------------------------------------------------------
    # 1) READ  (always)
    #    - Open the file with h5py and grab the dataset at data_path.
    #    - Read the whole cube when band is None, else a single slice
    #      dset[band]. Cast to float32 so the array can hold NaN.
    #    - Copy the attributes into a dict; pull out 'wavelengths' and
    #      'good_wavelengths' (a 0/1 flag per band).
    #    - Read the nodata_path mask as a boolean array.
    #
    # 2) CLEAN  (only when clean=True)
    #    Full cube (band is None) — exactly what you just did by hand:
    #      - keep only the good bands (boolean-index the band axis)
    #      - trim 'wavelengths' the same way
    #      - set the no-data pixels to NaN on every band
    #    Single band (band is given) — you hold just one 2-D slice:
    #      - if this band is NOT good, there is nothing to keep -> fill
    #        the whole slice with NaN
    #      - if it IS good, just NaN-out the no-data pixels
    #      - keep 'wavelengths'/'good_wavelengths' as 1-element arrays so
    #        the outputs still line up with the full-cube case
    #
    # 3) RETURN
    #    - Write the updated 'wavelengths' and 'good_wavelengths' back into
    #      attrs, then return  data, attrs, no_data

    # TODO: upgrade your simple loader into this cleaning version
    pass