"""
RockGlacierNet — compute per-band normalization stats for a built patch set.

The training data_generator z-score normalizes each band. Its built-in
GLOBAL_MEAN/STD were measured on the OLD glacier-centric crops and badly misfit the
whole-Sikkim region patches (e.g. Elev std 282 vs actual ~1150, NDVI off by +3.7 sigma).
This writes correct stats next to the patches; data_generator loads them automatically.

    python patch_dataset/compute_norm_stats.py --patch-size 128

Output: patch_dataset/patches_{size}/metadata/norm_stats.json
"""

import os
import sys
import json
import glob
import argparse
import numpy as np
import rasterio

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import config as C


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--patch-size", type=int, default=C.DEFAULT_PATCH)
    args = ap.parse_args()

    outdir = C.patches_dir(args.patch_size)
    feats = sorted(glob.glob(str(outdir / "train" / "features" / "*.tif")))
    if not feats:
        sys.exit(f"No train features in {outdir}/train/features — build patches first.")

    n = C.NUM_BANDS
    # Welford-ish running sums over all pixels of all train patches.
    count = 0
    s1 = np.zeros(n, dtype=np.float64)
    s2 = np.zeros(n, dtype=np.float64)
    for fp in feats:
        with rasterio.open(fp) as s:
            a = s.read().astype(np.float64)                 # (C, H, W)
        a = np.where(np.isfinite(a) & (a != C.NODATA_VALUE), a, 0.0)
        flat = a.reshape(n, -1)
        s1 += flat.sum(axis=1)
        s2 += (flat * flat).sum(axis=1)
        count += flat.shape[1]

    mean = s1 / count
    std = np.sqrt(np.maximum(s2 / count - mean * mean, 1e-12))

    stats = {"band_names": C.BAND_NAMES,
             "mean": [round(float(x), 6) for x in mean],
             "std":  [round(float(x), 6) for x in std],
             "num_train_patches": len(feats),
             "num_pixels": int(count)}
    out = outdir / "metadata" / "norm_stats.json"
    with open(out, "w") as f:
        json.dump(stats, f, indent=2)

    print(f"Computed over {len(feats)} train patches ({count:,} px/band)")
    print(f"{'band':7s} {'mean':>12s} {'std':>12s}")
    for nm, m, sd in zip(C.BAND_NAMES, stats["mean"], stats["std"]):
        print(f"{nm:7s} {m:12.4f} {sd:12.4f}")
    print(f"\nSaved -> {out}")


if __name__ == "__main__":
    main()