"""Re-score a Blue Marble Burn model on HLS Burn Scars (official validation split; the dataset has no separate test split).

    pip install numpy onnxruntime rasterio huggingface_hub
    python evaluate.py --model small/model_fp32.onnx           # or small/model_int8.onnx

Data: HLS Burn Scars (6-band HLS surface reflectance), about 2.6 GB, downloaded once. The six bands are placed at
B02, B03, B04, B8A, B11 and B12; the other bands are 0 after normalisation."""
import numpy as np

# Sentinel-2 L1C, 13 bands (B01, B02, B03, B04, B05, B06, B07, B08, B8A, B09, B10, B11, B12)
MEAN = np.array([2357.089, 2137.385, 2018.788, 2082.986, 2295.651, 2854.537, 3122.849, 3040.560, 3306.481, 1473.847, 506.070, 2472.825, 1838.929], np.float32)
STD = np.array([1624.683, 1675.806, 1557.708, 1833.702, 1823.738, 1733.977, 1732.131, 1679.732, 1727.26, 1024.687, 442.165, 1331.411, 1160.419], np.float32)


def normalize(dn, bands=None):
    """dn: [C, H, W] digital numbers for the 13 bands, or for the band indices in `bands`. Missing bands are 0 after normalisation."""
    out = np.zeros((13, *dn.shape[1:]), np.float32)
    idx = list(range(13)) if bands is None else list(bands)
    out[idx] = (dn.astype(np.float32) - MEAN[idx, None, None]) / STD[idx, None, None]
    return out


def session(path, threads=0):
    import onnxruntime as ort
    so = ort.SessionOptions()
    if threads: so.intra_op_num_threads = threads
    prov = [p for p in ("CUDAExecutionProvider", "CPUExecutionProvider") if p in ort.get_available_providers()]
    return ort.InferenceSession(path, so, providers=prov)


def predict(sess, x):
    """Whole image in one pass: pad H and W to multiples of 16 by edge replication, crop the logits back."""
    h, w = x.shape[-2:]
    xp = np.pad(x, ((0, 0), (0, (-h) % 16), (0, (-w) % 16)), mode="edge")[None].astype(np.float32)
    return sess.run(None, {sess.get_inputs()[0].name: xp})[0][0, :, :h, :w].argmax(0)


class MIoU:
    def __init__(self, classes):
        self.classes, n = classes, len(classes)
        self.cm = np.zeros((n, n), np.int64)

    def add(self, y, p):
        n = len(self.classes); m = y >= 0
        self.cm += np.bincount(n * y[m].astype(np.int64) + p[m], minlength=n * n).reshape(n, n)

    def report(self, title):
        inter = np.diag(self.cm).astype(float); union = self.cm.sum(0) + self.cm.sum(1) - inter
        iou = inter / np.maximum(union, 1)
        print(f"{title}: mIoU {100 * iou.mean():.2f}  (" + ", ".join(f"{c} {100 * v:.2f}" for c, v in zip(self.classes, iou)) + ")", flush=True)
        return float(iou.mean())


import argparse, glob, os, tarfile
import numpy as np, rasterio

BANDS = [1, 2, 3, 8, 11, 12]


def fetch(data):
    root = f"{data}/hls_burn_scars"
    if os.path.exists(f"{root}/.ok"): return root
    from huggingface_hub import hf_hub_download
    arc = f"{data}/hls_burn_scars.tar.gz"
    if not os.path.exists(arc): arc = hf_hub_download("ibm-nasa-geospatial/hls_burn_scars", "hls_burn_scars.tar.gz", repo_type="dataset", local_dir=data)
    os.makedirs(f"{root}/validation", exist_ok=True)
    with tarfile.open(arc, "r|gz") as t:
        for m in t:
            if m.isfile() and "validation/" in m.name and m.name.endswith(".tif"):
                m.name = os.path.basename(m.name); t.extract(m, f"{root}/validation", filter="data")
    open(f"{root}/.ok", "w").close()
    return root


def read(p):
    with rasterio.open(p) as s: return s.read()


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--model", required=True); ap.add_argument("--data", default="data")
    a = ap.parse_args(); os.makedirs(a.data, exist_ok=True)
    root = fetch(a.data); sess = session(a.model)
    fs = sorted(glob.glob(f"{root}/validation/*_merged.tif")); m = MIoU(["not burned", "burn scar"])
    for f in fs:
        dn = np.clip(np.nan_to_num(read(f).astype(np.float32)) * 10000, 0, 65535).astype(np.uint16)  # reflectance -> DN
        y = read(f.replace("_merged.tif", ".mask.tif"))[0].astype(np.int8).astype(np.int64)
        m.add(y, predict(sess, normalize(dn, BANDS)))
    m.report(f"HLS Burn Scars validation ({len(fs)} scenes)")


if __name__ == "__main__": main()