"""
Compute per-band normalization statistics (mean/std) over the TRAIN split of a
dataset and write them to <dataset>/metadata/norm_stats.json.

Stats are accumulated over VALID pixels only (NoData -9999 and non-finite values
are excluded). The RockGlacierDataGenerator auto-loads this file via its
_load_norm_stats() (it looks at <features_dir>/../../metadata/norm_stats.json),
so once this file exists the next training run uses the correct stats.

Usage:
    python compute_norm_stats.py                         # defaults to output/dataset
    python compute_norm_stats.py <train_features_dir>    # custom
"""
import os
import sys
import glob
import json
import numpy as np
import rasterio

BAND_NAMES = ["B2", "B3", "B4", "B8", "B11", "B12",
              "NDVI", "NDWI", "NDSI", "Elev", "Slope", "Aspect"]

DEFAULT_FEAT_DIR = "datacollection/output/dataset/train/features"


def main(feat_dir=DEFAULT_FEAT_DIR):
    files = sorted(glob.glob(os.path.join(feat_dir, "*.tif")))
    if not files:
        files = sorted(glob.glob(os.path.join(feat_dir, "*.npy")))
    assert files, f"No .tif/.npy features in {feat_dir}"

    nb = None
    ssum = ssq = cnt = None
    for fp in files:
        if fp.endswith(".npy"):
            arr = np.load(fp).astype(np.float64)            # (H, W, C)
        else:
            with rasterio.open(fp) as s:
                arr = np.moveaxis(s.read(), 0, -1).astype(np.float64)  # (H, W, C)
        if nb is None:
            nb = arr.shape[-1]
            ssum = np.zeros(nb); ssq = np.zeros(nb); cnt = np.zeros(nb)
        flat = arr.reshape(-1, nb)
        valid = np.isfinite(flat) & (flat != -9999)
        x = np.where(valid, flat, 0.0)
        ssum += x.sum(axis=0)
        ssq += (x * x).sum(axis=0)
        cnt += valid.sum(axis=0)

    mean = ssum / np.maximum(cnt, 1)
    var = ssq / np.maximum(cnt, 1) - mean ** 2
    std = np.sqrt(np.maximum(var, 1e-12))

    names = BAND_NAMES[:nb] if nb <= len(BAND_NAMES) else [f"band_{i}" for i in range(nb)]
    out = {
        "band_names": names,
        "mean": [float(m) for m in mean],
        "std": [float(s) for s in std],
        "num_train_images": len(files),
        "num_pixels": int(cnt.max()),
    }

    # <feat_dir>/../../metadata/norm_stats.json   (e.g. output/dataset/metadata/)
    meta_dir = os.path.normpath(os.path.join(feat_dir, os.pardir, os.pardir, "metadata"))
    os.makedirs(meta_dir, exist_ok=True)
    out_path = os.path.join(meta_dir, "norm_stats.json")
    with open(out_path, "w") as f:
        json.dump(out, f, indent=2)

    print(f"Computed norm stats from {len(files)} images, {int(cnt.max()):,} valid pixels/band")
    print(f"  {'band':6s} {'mean':>12s} {'std':>12s}")
    for n, m, s in zip(names, mean, std):
        print(f"  {n:6s} {m:12.4f} {s:12.4f}")
    print(f"\nWrote {out_path}")


if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else DEFAULT_FEAT_DIR)