#!/usr/bin/env python3
"""Small helpers for functional.sh (numpy + rasterio).

    helpers.py ndvi OUT.tif            write the 600x400 EPSG:4326 Float32 NDVI test raster
    helpers.py rinfo FILE.tif          print raster facts as JSON (driver, size, crs, dtype,
                                       nodata, layout, compression, overviews)
    helpers.py durations RUN.json      print "total=<s> nodes=<min>..<max>ms" for a run
    helpers.py elapsed DATASET.json    print updated_at - created_at of a dataset, seconds
"""
import json
import sys
from datetime import datetime


def ndvi(out):
    import numpy as np
    import rasterio
    from rasterio.transform import from_origin

    w, h = 600, 400  # 0.01 deg pixels, bbox [68.5, 51.5, 74.5, 55.5]
    y, x = np.mgrid[0:h, 0:w].astype(np.float32)
    v = 0.45 + 0.35 * np.sin(x / 47.0) * np.cos(y / 31.0) + 0.1 * np.sin((x + y) / 13.0)
    v = np.clip(v, -0.2, 1.0).astype(np.float32)
    v[:5, :] = -9999
    profile = dict(driver="GTiff", width=w, height=h, count=1, dtype="float32", crs="EPSG:4326",
                   transform=from_origin(68.5, 55.5, 0.01, 0.01), nodata=-9999)
    with rasterio.open(out, "w", **profile) as dst:
        dst.write(v, 1)


def rinfo(path):
    import rasterio

    with rasterio.open(path) as d:
        st = d.tags(ns="IMAGE_STRUCTURE")
        print(json.dumps({
            "driver": d.driver, "width": d.width, "height": d.height,
            "epsg": d.crs.to_epsg() if d.crs else None, "dtype": d.dtypes[0], "nodata": d.nodata,
            "layout": st.get("LAYOUT"), "compression": st.get("COMPRESSION"),
            "overviews": len(d.overviews(1)), "bounds": [round(b, 6) for b in d.bounds],
        }))


def ts(s):
    if not s:
        return None
    s = s.replace("Z", "+00:00")
    if "." in s:  # trim nanoseconds to microseconds for fromisoformat
        head, rest = s.split(".", 1)
        frac = "".join(ch for ch in rest if ch.isdigit())
        tz = rest[len(frac):]
        s = "%s.%s%s" % (head, frac[:6].ljust(6, "0"), tz)
    return datetime.fromisoformat(s)


def durations(path):
    with open(path, encoding="utf-8") as f:
        run = json.load(f)
    total = None
    if run.get("started_at") and run.get("finished_at"):
        total = (ts(run["finished_at"]) - ts(run["started_at"])).total_seconds()
    nd = [(ts(n["finished_at"]) - ts(n["started_at"])).total_seconds() * 1000
          for n in run.get("nodes", []) if n.get("started_at") and n.get("finished_at")]
    out = "total=%.2fs" % total if total is not None else "total=?"
    if nd:
        out += " nodes=%d..%dms" % (min(nd), max(nd))
    print(out)


def elapsed(path):
    with open(path, encoding="utf-8") as f:
        d = json.load(f)
    print("%.2f" % (ts(d["updated_at"]) - ts(d["created_at"])).total_seconds())


if __name__ == "__main__":
    cmd, arg = sys.argv[1], sys.argv[2]
    {"ndvi": ndvi, "rinfo": rinfo, "durations": durations, "elapsed": elapsed}[cmd](arg)