#!/usr/bin/env python3
"""Replace static channels in a current FireWx-FM stack.

The preferred path uses an already-audited 5 km static cache from the corrected
training cache. A raw-raster fallback is available for rebuilding that cache.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

import numpy as np
import rasterio
from affine import Affine
from rasterio.enums import Resampling
from rasterio.vrt import WarpedVRT
from rasterio.warp import reproject


STATIC_LAYERS = [
    ("fuel_fbfm40", 12, "LANDFIRE_FUELS_2024/FBFM40/LF2024_FBFM40_CONUS/Tif/LF2024_FBFM40_CONUS.tif", Resampling.nearest),
    ("canopy_cover", 13, "LANDFIRE_FUELS_2024/CC/LF2024_CC_CONUS/Tif/LF2024_CC_CONUS.tif", Resampling.nearest),
    ("housing_density", 14, "WRC_Housing_Density/extracted/HUDen_CONUS.tif", Resampling.bilinear),
    ("population", 15, "LandScan_Global_2024/extracted/landscan-global-2024.tif", Resampling.bilinear),
]

STATIC_CACHE_NAME_TO_CHANNEL = {
    "fuel_fbfm40": 12,
    "canopy_cover": 13,
    "housing_density": 14,
    "population": 15,
}


def summarize(arr: np.ndarray, mask: np.ndarray | None = None) -> dict[str, Any]:
    vals = arr[mask] if mask is not None else arr.ravel()
    vals = vals[np.isfinite(vals)]
    if vals.size == 0:
        return {"cells": 0}
    return {
        "cells": int(vals.size),
        "mean": float(vals.mean()),
        "q50": float(np.quantile(vals, 0.50)),
        "q90": float(np.quantile(vals, 0.90)),
        "q99": float(np.quantile(vals, 0.99)),
        "max": float(vals.max()),
        "fraction_zero": float((vals == 0.0).mean()),
        "fraction_le_0p25": float((vals <= 0.25).mean()),
    }


def load_lower48_mask(
    reference_tif: Path,
    lower48_mask_npy: Path | None = None,
) -> tuple[np.ndarray, str]:
    """Load the Lower-48 mask or derive it from the reference stack."""
    if lower48_mask_npy is not None:
        lower48 = np.load(lower48_mask_npy).astype(bool)
        return lower48, str(lower48_mask_npy)

    with rasterio.open(reference_tif) as ref:
        # Prefer the dataset mask because the live stack is expected to carry the
        # CONUS support footprint directly on the reference grid.
        lower48 = ref.dataset_mask() > 0
    if not np.any(lower48):
        raise ValueError(
            "Could not derive a Lower-48 mask from the reference stack; "
            "supply --lower48-mask-npy explicitly."
        )
    return lower48.astype(bool), f"{reference_tif}::dataset_mask()"


def read_on_reference_grid(path: Path, ref: rasterio.DatasetReader, resampling: Resampling) -> tuple[np.ndarray, np.ndarray]:
    """Read a source raster directly on the reference 5 km grid.

    The raw LANDFIRE/WRC layers are 30 m CONUS rasters. Reading the whole
    overlapping source window before reprojection can require many GB of RAM.
    WarpedVRT lets GDAL sample only the requested 5 km output grid.
    """
    with rasterio.open(path) as src:
        vrt_kwargs: dict[str, Any] = {
            "crs": ref.crs,
            "transform": ref.transform,
            "width": ref.width,
            "height": ref.height,
            "resampling": resampling,
            "dtype": "float32",
        }
        if src.nodata is not None:
            vrt_kwargs["src_nodata"] = src.nodata
            vrt_kwargs["nodata"] = np.nan
        with WarpedVRT(src, **vrt_kwargs) as vrt:
            masked = vrt.read(1, masked=True)

    data = np.asarray(masked.filled(np.nan), dtype=np.float32)
    valid = (~np.ma.getmaskarray(masked)) & np.isfinite(data) & (data > -9000.0)
    clean = np.where(valid, data, 0.0).astype(np.float32)
    return clean, valid.astype(np.float32)


def transform_from_centers(y: np.ndarray, x: np.ndarray) -> Affine:
    if y.ndim != 1 or x.ndim != 1 or len(y) < 2 or len(x) < 2:
        raise ValueError("Static cache lat/lon coordinate arrays must be 1-D with at least two cells.")
    dx = float(np.median(np.diff(x)))
    dy = float(np.median(np.abs(np.diff(y))))
    if dx <= 0 or dy <= 0:
        raise ValueError(f"Unexpected static cache grid spacing dx={dx}, dy={dy}")
    left = float(x[0]) - 0.5 * dx
    if y[0] > y[-1]:
        top = float(y[0]) + 0.5 * dy
        return Affine(dx, 0.0, left, 0.0, -dy, top)
    bottom = float(y[0]) - 0.5 * dy
    return Affine(dx, 0.0, left, 0.0, dy, bottom)


def replace_from_static_cache(stack: np.ndarray, static_cache_npz: Path, reference_tif: Path, lower48: np.ndarray) -> dict[str, Any]:
    print(f"[static] using corrected static cache: {static_cache_npz}", flush=True)
    with np.load(static_cache_npz, allow_pickle=True) as data:
        static = data["static"].astype(np.float32)
        static_valid = data["static_valid"].astype(np.float32)
        y_src = data["lat"].astype(np.float32)
        x_src = data["lon"].astype(np.float32)
        coord_crs = str(data["coord_crs"][0]) if "coord_crs" in data else "EPSG:5070"
        names = [str(v) for v in data["static_names"].tolist()]

    source_transform = transform_from_centers(y_src, x_src)
    static_aligned = np.zeros((static.shape[0], stack.shape[1], stack.shape[2]), dtype=np.float32)
    static_valid_aligned = np.zeros((1, stack.shape[1], stack.shape[2]), dtype=np.float32)
    with rasterio.open(reference_tif) as ref:
        dst_transform = ref.transform
        dst_crs = ref.crs
    print(
        f"[static] reproject cache grid {static.shape[1:]} -> reference grid {stack.shape[1:]}",
        flush=True,
    )
    reproject(
        source=static_valid[0],
        destination=static_valid_aligned[0],
        src_transform=source_transform,
        src_crs=coord_crs,
        dst_transform=dst_transform,
        dst_crs=dst_crs,
        src_nodata=None,
        dst_nodata=0.0,
        resampling=Resampling.nearest,
    )
    for i, name in enumerate(names):
        print(f"[static] align cache layer {name}", flush=True)
        reproject(
            source=static[i],
            destination=static_aligned[i],
            src_transform=source_transform,
            src_crs=coord_crs,
            dst_transform=dst_transform,
            dst_crs=dst_crs,
            src_nodata=None,
            dst_nodata=0.0,
            resampling=Resampling.nearest,
        )

    layer_summaries: dict[str, Any] = {}
    stack[11] = static_valid_aligned[0]
    for i, name in enumerate(names):
        channel = STATIC_CACHE_NAME_TO_CHANNEL.get(name)
        if channel is None:
            continue
        stack[channel] = static_aligned[i]
        layer_summaries[name] = {
            "channel": int(channel),
            "source_index": int(i),
            "path": str(static_cache_npz),
            "resampling": "cache_nearest_reproject_to_reference",
            "values_all": summarize(static_aligned[i]),
            "values_lower48": summarize(static_aligned[i], lower48),
        }
    return {
        "source": "corrected_static_cache",
        "static_cache_npz": str(static_cache_npz),
        "coord_crs": coord_crs,
        "source_shape": [int(v) for v in static.shape],
        "aligned_shape": [int(v) for v in static_aligned.shape],
        "source_names": names,
        "layer_summaries": layer_summaries,
    }


def write_stack_tif(path: Path, stack: np.ndarray, reference_tif: Path) -> None:
    with rasterio.open(reference_tif) as ref:
        profile = ref.profile.copy()
    profile.update(count=stack.shape[0], dtype="float32", nodata=None, compress="deflate")
    path.parent.mkdir(parents=True, exist_ok=True)
    with rasterio.open(path, "w", **profile) as dst:
        dst.write(stack.astype(np.float32))


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input-stack-npy", type=Path, required=True)
    parser.add_argument("--reference-stack-tif", type=Path, required=True)
    parser.add_argument("--raw-root", type=Path)
    parser.add_argument("--static-cache-npz", type=Path)
    parser.add_argument("--lower48-mask-npy", type=Path)
    parser.add_argument("--write-lower48-mask-npy", type=Path)
    parser.add_argument("--output-stack-npy", type=Path, required=True)
    parser.add_argument("--output-stack-tif", type=Path)
    parser.add_argument("--summary-json", type=Path, required=True)
    args = parser.parse_args()
    if args.static_cache_npz is None and args.raw_root is None:
        raise ValueError("Provide --static-cache-npz for the fast path or --raw-root for raw raster rebuild.")

    stack = np.load(args.input_stack_npy).astype(np.float32)
    if stack.ndim != 3 or stack.shape[0] != 16:
        raise ValueError(f"Expected [16,H,W] input stack, got {stack.shape}")
    lower48, lower48_source = load_lower48_mask(args.reference_stack_tif, args.lower48_mask_npy)
    if lower48.shape != stack.shape[1:]:
        raise ValueError(f"lower48 mask {lower48.shape} does not match stack {stack.shape[1:]}")
    if args.write_lower48_mask_npy is not None:
        args.write_lower48_mask_npy.parent.mkdir(parents=True, exist_ok=True)
        np.save(args.write_lower48_mask_npy, lower48.astype(bool))

    old_static = stack[11:16].copy()
    with rasterio.open(args.reference_stack_tif) as ref:
        if (ref.height, ref.width) != stack.shape[1:]:
            raise ValueError(f"reference grid {(ref.height, ref.width)} does not match stack {stack.shape[1:]}")

    if args.static_cache_npz is not None:
        replacement = replace_from_static_cache(stack, args.static_cache_npz, args.reference_stack_tif, lower48)
        layer_summaries = replacement["layer_summaries"]
    else:
        print("[static] rebuilding from raw rasters", flush=True)
        static_valids = []
        layer_summaries: dict[str, Any] = {}
        with rasterio.open(args.reference_stack_tif) as ref:
            for name, channel, rel_path, resampling in STATIC_LAYERS:
                print(f"[static] read {name}: {args.raw_root / rel_path}", flush=True)
                arr, valid = read_on_reference_grid(args.raw_root / rel_path, ref, resampling)
                stack[channel] = arr
                static_valids.append(valid)
                layer_summaries[name] = {
                    "channel": int(channel),
                    "path": str(args.raw_root / rel_path),
                    "resampling": str(resampling.name),
                    "valid_fraction_all": float((valid > 0.0).mean()),
                    "valid_fraction_lower48": float((valid[lower48] > 0.0).mean()),
                    "values_all": summarize(arr),
                    "values_lower48": summarize(arr, lower48),
                }
        static_valid = np.stack(static_valids, axis=0).mean(axis=0).astype(np.float32)
        stack[11] = static_valid
        replacement = {"source": "raw_rasters", "raw_root": str(args.raw_root)}

    args.output_stack_npy.parent.mkdir(parents=True, exist_ok=True)
    np.save(args.output_stack_npy, stack.astype(np.float32))
    if args.output_stack_tif is not None:
        write_stack_tif(args.output_stack_tif, stack, args.reference_stack_tif)

    summary = {
        "input_stack_npy": str(args.input_stack_npy),
        "reference_stack_tif": str(args.reference_stack_tif),
        "raw_root": str(args.raw_root) if args.raw_root is not None else None,
        "static_cache_npz": str(args.static_cache_npz) if args.static_cache_npz is not None else None,
        "lower48_mask_npy": str(args.lower48_mask_npy) if args.lower48_mask_npy is not None else None,
        "lower48_mask_source": lower48_source,
        "lower48_mask_written": str(args.write_lower48_mask_npy) if args.write_lower48_mask_npy is not None else None,
        "output_stack_npy": str(args.output_stack_npy),
        "output_stack_tif": str(args.output_stack_tif) if args.output_stack_tif is not None else None,
        "replacement": replacement,
        "shape": [int(v) for v in stack.shape],
        "old_static_valid_lower48": summarize(old_static[0], lower48),
        "new_static_valid_lower48": summarize(stack[11], lower48),
        "old_fuel_lower48": summarize(old_static[1], lower48),
        "new_fuel_lower48": summarize(stack[12], lower48),
        "layers": layer_summaries,
    }
    args.summary_json.parent.mkdir(parents=True, exist_ok=True)
    args.summary_json.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(json.dumps(summary, indent=2, sort_keys=True))
    return 0


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