#!/usr/bin/env python3
"""Write a FireWx-FM probability NPY to a georeferenced GeoTIFF."""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import rasterio


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--reference-stack", type=Path, required=True)
    parser.add_argument("--prediction-npy", type=Path, required=True)
    parser.add_argument("--output-tif", type=Path, required=True)
    args = parser.parse_args()

    pred = np.load(args.prediction_npy).astype(np.float32)
    with rasterio.open(args.reference_stack) as src:
        if (src.height, src.width) != pred.shape:
            raise ValueError(f"reference grid {(src.height, src.width)} does not match prediction {pred.shape}")
        profile = src.profile.copy()
    profile.update(count=1, dtype="float32", nodata=0.0, compress="deflate", predictor=2)

    args.output_tif.parent.mkdir(parents=True, exist_ok=True)
    with rasterio.open(args.output_tif, "w", **profile) as dst:
        dst.write(pred.astype(np.float32), 1)
    print(f"wrote {args.output_tif}")
    return 0


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