#!/usr/bin/env python3
"""
Mask alignment tool: crop and align a TIFF mask to a hyperspectral BSQ file,
then set areas outside the hyperspectral valid data mask to 0.
"""

import argparse
import numpy as np
import rasterio
from rasterio.warp import reproject, Resampling
from rasterio.enums import Resampling as ResamplingEnum
import os

def get_valid_mask_all_bands(hyperspectral_file):
    with rasterio.open(hyperspectral_file) as src:
        # Read all bands
        data = src.read()  # shape (bands, height, width)
        # Mask where all bands are zero
        all_zero = np.all(data == 0, axis=0)
        valid = ~all_zero
        return valid, src.profile

def align_mask_to_hyperspectral(mask_file, hyperspectral_profile, resampling=Resampling.bilinear):
    """
    Reproject/resample the mask to match the hyperspectral spatial grid.
    Returns the aligned mask as a 2D array.
    """
    with rasterio.open(mask_file) as mask_src:
        # Create an empty array for the output
        out_shape = (hyperspectral_profile['height'], hyperspectral_profile['width'])
        out_array = np.empty(out_shape, dtype=mask_src.dtypes[0])
        # Reproject the first band
        reproject(
            source=rasterio.band(mask_src, 1),
            destination=out_array,
            src_transform=mask_src.transform,
            src_crs=mask_src.crs,
            dst_transform=hyperspectral_profile['transform'],
            dst_crs=hyperspectral_profile['crs'],
            resampling=resampling
        )
        return out_array

def main():
    parser = argparse.ArgumentParser(
        description="Crop TIFF mask to hyperspectral extent and zero-out outside valid data mask."
    )
    parser.add_argument('hyperspectral', help='Input hyperspectral BSQ file (with .hdr)')
    parser.add_argument('mask', help='Input TIFF mask file')
    parser.add_argument('output', help='Output TIFF file')
    parser.add_argument('--band', type=int, default=0,
                        help='Band index (0-based) to use for valid data mask (default: 0)')
    parser.add_argument('--nodata', type=float, default=None,
                        help='NoData value in hyperspectral data; if not given, uses file metadata')
    parser.add_argument('--resample', choices=['nearest', 'bilinear', 'cubic'], default='bilinear',
                        help='Resampling method for mask alignment (default: bilinear)')
    args = parser.parse_args()

    # Mapping of resampling names to rasterio enums
    resample_map = {
        'nearest': Resampling.nearest,
        'bilinear': Resampling.bilinear,
        'cubic': Resampling.cubic
    }
    resampling = resample_map[args.resample]

    # Get valid data mask from hyperspectral file
    print(f"Reading valid mask from {args.hyperspectral} (band {args.band})...")
    valid_mask, hs_profile = get_valid_mask_all_bands(args.hyperspectral)
    print(f"Hyperspectral dimensions: {hs_profile['height']} x {hs_profile['width']}")

    # Reproject mask to hyperspectral grid
    print(f"Aligning mask {args.mask} to hyperspectral grid...")
    aligned_mask = align_mask_to_hyperspectral(args.mask, hs_profile, resampling)

    # Zero-out areas outside valid_mask
    print("Applying valid mask...")
    result = np.where(valid_mask, aligned_mask, 0).astype(aligned_mask.dtype)

    # Save result as TIFF
    print(f"Saving to {args.output}...")
    out_profile = hs_profile.copy()
    out_profile.update({
        'driver': 'GTiff',
        'count': 1,
        'dtype': result.dtype,
        'compress': 'lzw'          # optional compression
    })
    with rasterio.open(args.output, 'w', **out_profile) as dst:
        dst.write(result, 1)

    print("Done.")

if __name__ == '__main__':
    main()