#!/usr/bin/env python3
"""Downloads the LROC WAC Hapke-normalized 643 nm mosaic of the Moon and makes the
renderer's Moon image from it.

Source: LRO Camera Wide Angle Camera mosaic, 643 nm band, photometrically normalized
by Hapke parameter maps to incidence 60 deg, emission 0 deg, phase 60 deg (Sato et
al. 2014, JGR Planets 119, 1775; https://www.lroc.asu.edu/posts/801), 400 m per pixel
(76 pixels per degree), equirectangular, 70 S to 70 N, as eight PDS tiles of 32-bit
floats (product WAC_HAPKE_643NM_*, data set LRO-L-LROC-5-RDR-V1.0, Arizona State
University).

The tiles go to data/moon/source/ and are checked by sha256. They are joined into one
mosaic with longitude -180 to 180 left to right and latitude +90 at the top, like the
Earth image. Missing pixels, and the polar caps beyond 70 degrees that the mosaic does
not cover (6 % of the surface), get the mosaic's mean, weighted by area. The result is
resampled by area to 16384 x 8192, the largest texture lavapipe takes (666 m per pixel
at the equator), and written as data/moon/wac-hapke-643.png: 16-bit gray, each value
the normalized reflectance I/F times 100,000.

With --shape, it instead prints the Moon's mean reflectance in each of the WAC's seven
bands relative to 643 nm over the whole mosaic, weighted by area: it downloads every
tile of every band in turn (56 files, 8 GB) and keeps none of them. That is how the spectrum in crates/resolved/src/moon.rs
was found.

Needs numpy and Pillow. Building the image takes about 20 s and 2.6 GB of memory once
the tiles are downloaded (about a minute at 50 MB/s, with a pause between files).
"""

import hashlib
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

import numpy as np
from PIL import Image

BASE = (
    "https://pds.mcp.nasa.gov/data/store/img/lunar_reconnaissance_orbiter/pds4/lroc/"
    "lro-l-lroc-5-rdr/LROLRC_2001/DATA/MDR/WAC_HAPKE/"
)
PPD = 76  # pixels per degree
WIDTH = 6840  # samples per tile: 90 degrees
RECORD = WIDTH * 4  # bytes per line, and the label's one record
# Tile, its western longitude (0-360 E), north or south, and sha256 of the file.
TILES = [
    ("E350N0450", 0, "N",
     "1612ded2849aa784be4b1d7d8b46d6778e342a7838f5f59cd55373ac7347f2ec"),
    ("E350N1350", 90, "N",
     "1be241a613ed2141bd7a7b62d3b050a9e7d20b4e15aaae92fb665558bcf23601"),
    ("E350N2250", 180, "N",
     "8aa0b7b69501e606e3fd4968d99787443fd629bf7a2682b0854a8caecb759a1e"),
    ("E350N3150", 270, "N",
     "a92b6253b062b4c7b2a8da6eb0c8bb844b6d97c1bc5a518177970046c8e8b20a"),
    ("E350S0450", 0, "S",
     "967eb2a5fcbc7c550b1e411f7cb843c1d8c29f53d1f069353ec18fadaa05ef60"),
    ("E350S1350", 90, "S",
     "dc5c6b86f3a4dc6c6e3eb6e4766371b75aa97482980766377d7a3abb6775ad4c"),
    ("E350S2250", 180, "S",
     "27f487579d0543daa630961b1922125d03fed685871d5088c1d4b4a15c5b8578"),
    ("E350S3150", 270, "S",
     "4aeea79efcb6402b012abf07e41cbba9ca8c2e9af83535de67d1b3345f22918c"),
]
BANDS = [321, 360, 415, 566, 604, 643, 689]
# Lines per tile: north tiles have one more, just above 70 N, which is dropped.
LINES = {"N": 5321, "S": 5320}
ROWS = 70 * PPD  # 5320 lines from the equator to 70 degrees
OUT_W, OUT_H = 16384, 8192
SCALE = 100_000  # stored value per unit I/F
STRIP = 512  # rows or columns resampled at a time


def url(band, tile):
    return f"{BASE}WAC_HAPKE_{band}NM_{tile}.IMG"


def download(address):
    """The whole file at `address`. One request at a time, a pause between them, and
    when the server answers 429 (too many requests) the wait it asks for: the PDS
    server blocks a client for an hour after a burst of requests."""
    time.sleep(5)
    for _ in range(6):
        try:
            with urllib.request.urlopen(address) as response:
                return response.read()
        except urllib.error.HTTPError as e:
            if e.code != 429:
                raise
            wait = int(e.headers.get("Retry-After", "600"))
            print(f"  rate-limited; waiting {wait} s", flush=True)
            time.sleep(wait + 30)
    sys.exit(f"{address}: still rate-limited")


def fetch(band, tile, sha, directory):
    path = directory / f"WAC_HAPKE_{band}NM_{tile}.IMG"
    if not path.exists():
        print(f"fetching {path.name}", flush=True)
        part = path.with_suffix(".part")
        part.write_bytes(download(url(band, tile)))
        part.rename(path)
    digest = hashlib.sha256(path.read_bytes()).hexdigest()
    if sha is None:
        print(f"  sha256 {digest}  {path.name}  (not pinned)")
    elif digest != sha:
        sys.exit(f"{path}: sha256 {digest}, expected {sha}")
    return path


def open_tile(path, hemisphere):
    """The tile's 5320 lines from the equator outward, as rows from north to south,
    mapped from the file rather than read into memory."""
    data = np.memmap(path, dtype="<f4", mode="r", offset=RECORD)
    data = data.reshape(LINES[hemisphere], WIDTH)
    if hemisphere == "N":
        data = data[1:]  # the line centered just above 70 N
    return data


def mosaic_rows(tiles, start, stop):
    """Rows `start` to `stop` of the hemisphere `tiles` make side by side, as float32
    with the PDS special values (null, saturation; all below -3e38) as NaN."""
    rows = np.hstack([t[start:stop] for t in tiles]).astype(np.float32)
    rows[rows < -1e30] = np.nan
    return rows


def row_weights(rows, north):
    """cos(latitude) of each row's center, for area weighting."""
    k = np.arange(rows) + 0.5
    lat = (ROWS - k) / PPD if north else -k / PPD
    return np.cos(np.radians(lat))


def resample(image, axis, size):
    """Area-weighted resampling of `image` along `axis` to `size` samples."""
    image = np.moveaxis(image, axis, 0).astype(np.float64)
    n = image.shape[0]
    # Running integral: total[i] is the sum of samples below i.
    total = np.concatenate([np.zeros((1,) + image.shape[1:]), np.cumsum(image, axis=0)])

    def integral(t):
        i = np.minimum(np.floor(t).astype(int), n - 1)
        frac = (t - i)[:, None] if image.ndim > 1 else t - i
        return total[i] + frac * image[i]

    edges = np.arange(size + 1) * (n / size)
    out = (integral(edges[1:]) - integral(edges[:-1])) * (size / n)
    return np.moveaxis(out.astype(np.float32), 0, axis)


def build(root):
    """Works through the mosaic STRIP rows (or columns) at a time, so memory stays near
    the 0.9 GB of the half-resampled image instead of the many GB a whole-mosaic
    resampling in float64 takes."""
    source = root / "source"
    source.mkdir(parents=True, exist_ok=True)
    # Columns from -180: the tiles starting at 180 E and 270 E come first.
    order = sorted(TILES, key=lambda t: (t[1] + 180) % 360)
    halves = [
        ([open_tile(fetch(643, t, s, source), h) for t, w, h, s in order if h == hemisphere],
         row_weights(ROWS, hemisphere == "N"))
        for hemisphere in "NS"
    ]
    # The mean of the valid pixels, by area.
    total = weight = 0.0
    missing = 0
    for tiles, weights in halves:
        for start in range(0, ROWS, STRIP):
            rows = mosaic_rows(tiles, start, start + STRIP)
            w = weights[start:start + STRIP, None]
            valid = ~np.isnan(rows)
            total += float(np.nansum(rows * w))
            weight += float((w * valid).sum())
            missing += int((~valid).sum())
    mean = total / weight
    pixels = 2 * ROWS * 4 * WIDTH
    print(f"mosaic {4 * WIDTH} x {2 * ROWS}: mean I/F {mean:.5f}, "
          f"{missing} missing pixels ({100 * missing / pixels:.4f} %)")
    # Missing pixels and the polar caps get the mean. Resampled across, then down.
    cap = 20 * PPD
    height = 2 * cap + 2 * ROWS  # 90 N to 90 S
    print(f"resampling {4 * WIDTH} x {height} to {OUT_W} x {OUT_H}", flush=True)
    across = np.empty((height, OUT_W), dtype=np.float32)
    across[:cap] = mean
    across[height - cap:] = mean
    for k, (tiles, _) in enumerate(halves):
        for start in range(0, ROWS, STRIP):
            rows = mosaic_rows(tiles, start, start + STRIP)
            rows[np.isnan(rows)] = mean
            top = cap + k * ROWS + start
            across[top:top + len(rows)] = resample(rows, 1, OUT_W)
    values = np.empty((OUT_H, OUT_W), dtype=np.uint16)
    low, high = np.inf, -np.inf
    for start in range(0, OUT_W, STRIP):
        down = resample(across[:, start:start + STRIP], 0, OUT_H)
        low, high = min(low, float(down.min())), max(high, float(down.max()))
        values[:, start:start + STRIP] = np.clip(np.round(down * SCALE), 0, 65535)
    out = root / "wac-hapke-643.png"
    Image.fromarray(values).save(out)
    print(f"wrote {out}: I/F from {low:.4f} to {high:.4f}")


def band_mean(band):
    """Area-weighted mean I/F of all eight tiles of `band`, each downloaded, summed and
    discarded in turn."""
    total = weight = 0.0
    for tile, w, hemisphere, s in TILES:
        print(f"  {band} nm {tile}", flush=True)
        data = np.frombuffer(download(url(band, tile)), dtype="<f4", offset=RECORD)
        rows = data.reshape(LINES[hemisphere], WIDTH)
        if hemisphere == "N":
            rows = rows[1:]
        rows = rows.copy()
        rows[rows < -1e30] = np.nan
        cos = row_weights(ROWS, hemisphere == "N")[:, None]
        total += float(np.nansum(rows * cos))
        weight += float(np.sum(~np.isnan(rows) * cos))
    return total / weight


def shape():
    means = {band: band_mean(band) for band in BANDS}
    for band in BANDS:
        print(f"{band} nm: mean I/F {means[band]:.5f}, relative to 643 nm "
              f"{means[band] / means[643]:.4f}")


if __name__ == "__main__":
    root = Path(__file__).resolve().parent.parent / "data" / "moon"
    if "--shape" in sys.argv[1:]:
        shape()
    else:
        build(root)