#!/usr/bin/env python3
# Turns NASA SVS's "Deep Star Maps 2020" Milky Way background (the star map
# without its Hipparcos and Tycho stars, which the sky draws itself) into the
# sRGB JPEG the sky component maps onto the celestial sphere.
#
#   curl -O https://svs.gsfc.nasa.gov/vis/a000000/a004800/a004851/milkyway_2020_4k.exr
#   python3 tools/sky/build-milkyway.py milkyway_2020_4k.exr public/assets/sky/milkyway-4k.jpg
#
# Needs vips (to read the OpenEXR), numpy and Pillow. The map is plate carrée
# in J2000 right ascension and declination, centred on 0h with r.a. increasing
# to the left; this keeps that layout. Values stay linear in brightness up to
# WHITE (only the heart of the Sagittarius star cloud goes past it), and are
# stored with the sRGB curve so the faint parts keep their detail in 8 bits.
import os
import subprocess
import sys
import tempfile

import numpy as np
from PIL import Image

WHITE = 0.6


def main(src, dst):
    width, height = (int(subprocess.check_output(['vipsheader', '-f', f, src], stderr=subprocess.DEVNULL)) for f in ('width', 'height'))
    with tempfile.TemporaryDirectory() as tmp:
        raw = os.path.join(tmp, 'map.raw')
        subprocess.run(['vips', 'rawsave', src, raw], check=True, stderr=subprocess.DEVNULL)
        rgb = np.fromfile(raw, dtype=np.float32).reshape(height, width, -1)[..., :3]
    x = np.clip(rgb / WHITE, 0.0, 1.0)
    srgb = np.where(x <= 0.0031308, 12.92 * x, 1.055 * np.power(x, 1 / 2.4) - 0.055)
    # A little noise before rounding keeps the dark gradients from banding.
    noise = np.random.default_rng(1).uniform(-0.5, 0.5, srgb.shape)
    out = np.clip(np.round(srgb * 255 + noise), 0, 255).astype(np.uint8)
    Image.fromarray(out, 'RGB').save(dst, quality=88, optimize=True, progressive=True)
    print(f'{width}x{height} -> {dst} ({os.path.getsize(dst) // 1024} KB)')


if __name__ == '__main__':
    main(sys.argv[1], sys.argv[2])