"""Compare two rasters for shape, nodata patterns, and value agreement."""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import rasterio


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Compare two rasters used by the forest-canopy example."
    )
    parser.add_argument("--raster-a", required=True, type=Path, help="First raster path.")
    parser.add_argument("--raster-b", required=True, type=Path, help="Second raster path.")
    parser.add_argument(
        "--nodata-value",
        action="append",
        type=float,
        default=[],
        help="Additional nodata fill value to count explicitly. Repeat as needed.",
    )
    return parser


def _read_raster(path: Path) -> tuple[np.ndarray, float | None]:
    with rasterio.open(path) as src:
        return src.read(1), src.nodata


def _count_special_values(array: np.ndarray, nodata_values: list[float]) -> dict[str, int]:
    stats = {"nan": int(np.count_nonzero(np.isnan(array)))}
    for value in nodata_values:
        stats[str(value)] = int(np.count_nonzero(array == value))
    return stats


def compare_rasters(raster_a: Path, raster_b: Path, extra_nodata_values: list[float]) -> int:
    array_a, nodata_a = _read_raster(raster_a)
    array_b, nodata_b = _read_raster(raster_b)

    nodata_values: list[float] = []
    for candidate in [nodata_a, nodata_b, *extra_nodata_values]:
        if candidate is not None and candidate not in nodata_values:
            nodata_values.append(candidate)

    print(f"raster_a: {raster_a}")
    print(f"raster_b: {raster_b}")
    print(f"shape_a: {array_a.shape}")
    print(f"shape_b: {array_b.shape}")

    shape_match = array_a.shape == array_b.shape
    print(f"same_shape: {shape_match}")

    stats_a = _count_special_values(array_a, nodata_values)
    stats_b = _count_special_values(array_b, nodata_values)
    print(f"special_counts_a: {stats_a}")
    print(f"special_counts_b: {stats_b}")

    if shape_match:
        print(f"same_nodata_mask: {np.array_equal(np.isnan(array_a), np.isnan(array_b))}")
        comparable_mask = ~np.isnan(array_a) & ~np.isnan(array_b)
        print(f"same_values_where_not_nan: {np.array_equal(array_a[comparable_mask], array_b[comparable_mask])}")
    else:
        print("same_nodata_mask: False")
        print("same_values_where_not_nan: False")

    return 0 if shape_match else 1


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    return compare_rasters(args.raster_a, args.raster_b, args.nodata_value)


if __name__ == "__main__":
    raise SystemExit(main())
