"""Post-process seaweed masks with no-data and conservative water constraints."""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import rasterio
from rasterio.windows import Window


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--image", required=True, help="Source 4-band fused/multispectral raster.")
    parser.add_argument("--mask", required=True, help="Predicted binary mask raster.")
    parser.add_argument("--output", required=True, help="Filtered output mask raster.")
    parser.add_argument("--stripe-height", type=int, default=1024)
    parser.add_argument("--black-threshold", type=float, default=32.0, help="Pixels with all bands <= this are no-data.")
    parser.add_argument(
        "--land-ndwi-threshold",
        type=float,
        default=-0.35,
        help="Very conservative NDWI cutoff for obvious dry land. Lower is safer for floating algae.",
    )
    parser.add_argument(
        "--land-nir-ratio",
        type=float,
        default=1.6,
        help="Only suppress NDWI-low pixels when NIR is this many times brighter than green.",
    )
    parser.add_argument("--no-land-filter", action="store_true", help="Only remove black/no-data regions.")
    return parser.parse_args()


def obvious_land_mask(tile: np.ndarray, ndwi_threshold: float, nir_ratio: float) -> np.ndarray:
    """Return a conservative land mask from B,G,R,NIR-like 4-band data.

    This is not a substitute for an official coastline/water mask. It only removes
    strongly land-like pixels to avoid deleting real floating algae.
    """
    if tile.shape[0] < 4:
        return np.zeros(tile.shape[1:], dtype=bool)
    green = tile[1].astype(np.float32, copy=False)
    nir = tile[3].astype(np.float32, copy=False)
    ndwi = (green - nir) / (green + nir + 1e-6)
    return (ndwi < ndwi_threshold) & (nir > green * nir_ratio)


def main() -> None:
    args = parse_args()
    image_path = Path(args.image)
    mask_path = Path(args.mask)
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    with rasterio.open(image_path) as image_src, rasterio.open(mask_path) as mask_src:
        if (image_src.width, image_src.height) != (mask_src.width, mask_src.height):
            raise ValueError(
                f"Image and mask sizes differ: image={image_src.width}x{image_src.height}, "
                f"mask={mask_src.width}x{mask_src.height}"
            )
        profile = mask_src.profile.copy()
        profile.update(count=1, dtype="uint8", compress="lzw", nodata=0)

        total_pixels = image_src.width * image_src.height
        input_fg = 0
        output_fg = 0
        invalid_pixels = 0
        land_pixels = 0

        with rasterio.open(output_path, "w", **profile) as dst:
            for y in range(0, image_src.height, args.stripe_height):
                height = min(args.stripe_height, image_src.height - y)
                window = Window(0, y, image_src.width, height)
                image = image_src.read(window=window)
                mask = mask_src.read(1, window=window)

                predicted = mask > 0
                valid = np.max(image, axis=0) > args.black_threshold
                land = np.zeros(valid.shape, dtype=bool)
                if not args.no_land_filter:
                    land = obvious_land_mask(image, args.land_ndwi_threshold, args.land_nir_ratio)

                filtered = predicted & valid & ~land
                dst.write((filtered.astype(np.uint8) * 255), 1, window=window)

                input_fg += int(predicted.sum())
                output_fg += int(filtered.sum())
                invalid_pixels += int((~valid).sum())
                land_pixels += int(land.sum())

    print(f"image={image_path}")
    print(f"mask={mask_path}")
    print(f"output={output_path}")
    print(f"total_pixels={total_pixels}")
    print(f"input_foreground={input_fg} ratio={input_fg / total_pixels:.6f}")
    print(f"output_foreground={output_fg} ratio={output_fg / total_pixels:.6f}")
    print(f"removed_foreground={input_fg - output_fg} ratio={(input_fg - output_fg) / total_pixels:.6f}")
    print(f"invalid_or_black_pixels={invalid_pixels} ratio={invalid_pixels / total_pixels:.6f}")
    print(f"conservative_land_pixels={land_pixels} ratio={land_pixels / total_pixels:.6f}")


if __name__ == "__main__":
    main()