"""Package the existing SDR PNG and matching SVG as a backwards-compatible JPEG.

Run: uv run python avatar_ultrahdr.py
"""

import argparse
import io
from pathlib import Path

import numpy as np
from PIL import Image
from pyultrahdr import build_ultra_hdr

from avatar_color import BT2020_Y, SRGB_Y, validate_peak
from avatar_hdr import render_linear_hdr


def export_ultrahdr(png_path, svg_path, output, peak_nits=400):
    validate_peak(peak_nits)
    with Image.open(png_path) as image:
        sdr = np.asarray(image.convert('RGB'))
    hdr = render_linear_hdr(Path(svg_path).read_text(encoding='utf-8'), peak_nits)
    if hdr.shape != sdr.shape:
        raise ValueError('The SDR PNG and SVG must have matching dimensions and content')
    encoded = sdr.astype(np.float64) / 255
    linear = np.where(encoded <= .04045, encoded / 12.92, ((encoded + .055) / 1.055) ** 2.4)
    sdr_white = 203.0
    sdr_y = linear @ np.asarray(SRGB_Y) * sdr_white
    hdr_y = hdr @ np.asarray(BT2020_Y)
    # A single-channel map scales RGB together, preserving the base chromaticity.
    # Keep SDR shadows rather than forcing the PQ version's darker black level.
    gain = np.log2(np.maximum(hdr_y / np.maximum(sdr_y, 1e-8), 1))
    gain_max = max(float(gain.max()), 1e-6)
    # The library's capacity argument is a reference-white brightness, not our
    # colored pixels' luminance. Full gain must wait for the full gain headroom.
    result = build_ultra_hdr(sdr, gain, 0.0, gain_max,
                             sdr_white * 2 ** gain_max, sdr_white,
                             base_quality=100, gainmap_quality=100)
    Path(output).write_bytes(result)
    return verify_ultrahdr(result, sdr)


def verify_ultrahdr(result, sdr):
    with Image.open(io.BytesIO(result)) as image:
        decoded = np.asarray(image.convert('RGB'))
        directory = image._getmp()
        frames = directory[45057]
        entry = directory[45058][1]
        start = image.info['mpoffset'] + entry['DataOffset']
        with Image.open(io.BytesIO(result[start:start + entry['Size']])) as gain_image:
            gain_image.load()
            if gain_image.size != image.size or b'GainMapMax' not in gain_image.info['xmp']:
                raise ValueError('Missing full-resolution gain map or HDR metadata')
    error = np.abs(decoded.astype(float) - sdr)
    if frames != 2:
        raise ValueError('Ultra HDR output must contain a base and gain-map image')
    return {'size': list(sdr.shape[:2]), 'jpeg_frames': frames,
            'sdr_mean_absolute_error': float(error.mean()),
            'sdr_max_error': float(error.max())}


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--png', type=Path, default=Path('avatar.png'))
    parser.add_argument('--svg', type=Path, default=Path('avatar.svg'))
    parser.add_argument('--output', type=Path, default=Path('avatar-ultrahdr.jpg'))
    parser.add_argument('--peak-nits', type=float, default=400)
    args = parser.parse_args()
    print(export_ultrahdr(args.png, args.svg, args.output, args.peak_nits))
    print(f'Saved {args.output.resolve()}')