make_prediction.py

"""
VIIRS-frequency fire probability map — no ML training required.

Builds the probability map at 300 m (VIIRS native ~375 m), then warps
to 30 m on disk via streaming — never allocates the full 30 m grid in RAM.

Usage:
    python make_prediction.py
"""

import numpy as np
import pandas as pd
import rasterio
from rasterio.transform import from_bounds
from rasterio.warp import reproject, Resampling
import rasterio.crs
from scipy.ndimage import gaussian_filter
from pyproj import Transformer
from pathlib import Path
from datetime import datetime
import yaml


COARSE_FACTOR = 10   # work at 300 m; upsample to 30 m at the end


def main():
    with open("config/config.yaml") as f:
        cfg = yaml.safe_load(f)

    proc_dir  = Path(cfg["data"]["processed_dir"])
    pred_dir  = Path(cfg["output"]["prediction_dir"])
    pred_dir.mkdir(parents=True, exist_ok=True)
    raw_viirs = Path(cfg["data"]["raw_dir"]) / "viirs"

    # Reference grid from the aligned VIIRS raster (30 m)
    ref_path = proc_dir / "viirs_labels_aligned.tif"
    with rasterio.open(ref_path) as src:
        ref_transform = src.transform
        crs            = src.crs
        full_h, full_w = src.height, src.width
        ref_profile    = src.profile.copy()

    # Coarse grid (300 m)
    coarse_h = full_h // COARSE_FACTOR
    coarse_w = full_w // COARSE_FACTOR
    left   = ref_transform.c
    top    = ref_transform.f
    right  = left + ref_transform.a * full_w
    bottom = top  + ref_transform.e * full_h
    coarse_transform = from_bounds(left, bottom, right, top, coarse_w, coarse_h)

    print(f"Full grid : {full_h}×{full_w} px @ 30 m")
    print(f"Coarse grid: {coarse_h}×{coarse_w} px @ 300 m  "
          f"({coarse_h*coarse_w*4/1e6:.1f} MB per array)")

    years = cfg["data"]["viirs"]["years"]
    transformer = Transformer.from_crs("EPSG:4326", crs.to_epsg(), always_xy=True)

    yearly_burns  = np.zeros((coarse_h, coarse_w), dtype=np.float32)
    n_valid_years = 0

    for year in years:
        csv_path = raw_viirs / f"viirs_{year}_filtered.csv"
        if not csv_path.exists():
            print(f"  {year}: CSV not found, skipping")
            continue
        df = pd.read_csv(csv_path)
        if df.empty:
            print(f"  {year}: empty, skipping")
            continue

        xs, ys = transformer.transform(df["longitude"].values, df["latitude"].values)
        rows, cols = rasterio.transform.rowcol(coarse_transform, xs, ys)
        rows = np.asarray(rows)
        cols = np.asarray(cols)
        valid = (rows >= 0) & (rows < coarse_h) & (cols >= 0) & (cols < coarse_w)
        mask  = np.zeros((coarse_h, coarse_w), dtype=np.float32)
        mask[rows[valid], cols[valid]] = 1.0
        yearly_burns  += mask
        n_valid_years += 1
        print(f"  {year}: {valid.sum():,} detections -> {int(mask.sum()):,} coarse pixels burned")

    if n_valid_years == 0:
        print("No per-year CSVs found — using aggregated viirs_labels_aligned.tif")
        with rasterio.open(ref_path) as src:
            from rasterio.enums import Resampling as RsEnum
            raw_coarse = src.read(
                1,
                out_shape=(coarse_h, coarse_w),
                resampling=RsEnum.max
            ).astype(np.float32)
        raw_coarse[raw_coarse == 255] = 0
        yearly_burns  = raw_coarse
        n_valid_years = 1

    # Fire frequency
    freq = yearly_burns / n_valid_years

    # Gaussian smoothing — sigma=1.2 pixels at 300 m ≈ 360 m spread
    prob_coarse = gaussian_filter(freq, sigma=1.2).astype(np.float32)
    pmax = prob_coarse.max()
    if pmax > 0:
        prob_coarse /= pmax

    # Class map at coarse resolution
    thr        = cfg["model"]["thresholds"]
    cls_coarse = np.zeros((coarse_h, coarse_w), dtype=np.uint8)
    cls_coarse[prob_coarse > thr["low"]]      = 1
    cls_coarse[prob_coarse > thr["moderate"]] = 2
    cls_coarse[prob_coarse > thr["high"]]     = 3

    print(f"\nCoarse risk summary:")
    for c, name in [(3,"High"),(2,"Moderate"),(1,"Low"),(0,"Nil")]:
        n = int((cls_coarse==c).sum())
        print(f"  {name:10s}: {n:>8,} coarse pixels")

    # Upsample + save directly to 30 m TIFs via streaming reproject
    date_str  = datetime.now().strftime("%Y-%m-%d")
    prob_path = pred_dir / f"fire_probability_{date_str}.tif"
    cls_path  = pred_dir / f"fire_class_{date_str}.tif"

    out_profile = ref_profile.copy()
    out_profile.update(dtype="float32", count=1, nodata=None,
                       compress="lzw", tiled=True, blockxsize=256, blockysize=256)

    print(f"\nUpsampling to 30 m and saving …")

    with rasterio.open(prob_path, "w", **out_profile) as dst:
        reproject(
            source=prob_coarse[np.newaxis],
            destination=rasterio.band(dst, 1),
            src_transform=coarse_transform,
            src_crs=crs,
            dst_transform=ref_transform,
            dst_crs=crs,
            resampling=Resampling.bilinear,
        )

    cls_profile = out_profile.copy()
    cls_profile.update(dtype="uint8")
    with rasterio.open(cls_path, "w", **cls_profile) as dst:
        reproject(
            source=cls_coarse[np.newaxis].astype(np.uint8),
            destination=rasterio.band(dst, 1),
            src_transform=coarse_transform,
            src_crs=crs,
            dst_transform=ref_transform,
            dst_crs=crs,
            resampling=Resampling.nearest,
        )

    print(f"Probability raster -> {prob_path}")
    print(f"Class raster       -> {cls_path}")
    print("\nDone. Run next: python main.py --step simulate")


if __name__ == "__main__":
    main()