#!/usr/bin/env python3
"""Upscales every PNG in one directory into another, via one ONNX Runtime
CPU session shared across the whole batch (same reasoning as mdlp's own
Go-side batching: the fixed per-invocation model-load cost otherwise
dominates for the many small sprite frames one upscale definition has).

Runs on CPU only, not Vulkan: this Docker image has no real GPU, and
Vulkan software emulation (Mesa lavapipe) gives wrong colors for ML
inference workloads.

Usage: upscale_batch.py <model.onnx> <in-dir> <out-dir> <scale>
"""
import sys
import os

import numpy as np
import onnxruntime as ort
from PIL import Image


def main():
    model_path, in_dir, out_dir, scale_arg = sys.argv[1:5]
    scale = int(scale_arg)

    session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
    input_name = session.get_inputs()[0].name

    names = sorted(n for n in os.listdir(in_dir) if n.lower().endswith(".png"))
    for name in names:
        img = Image.open(os.path.join(in_dir, name)).convert("RGBA")
        rgb = img.convert("RGB")
        alpha = img.split()[3]

        # The model has no concept of transparency (3-channel RGB in/out) -
        # run it on the flattened RGB only, and scale the alpha mask
        # separately with a plain resize. Alpha edges don't need AI
        # enhancement, just to end up the right size.
        arr = np.asarray(rgb, dtype=np.float32) / 255.0
        arr = np.transpose(arr, (2, 0, 1))[np.newaxis, ...]
        output = session.run(None, {input_name: arr})[0]
        output = np.clip(output[0], 0.0, 1.0) * 255.0
        output = np.transpose(output, (1, 2, 0)).astype(np.uint8)
        upscaled_rgb = Image.fromarray(output, "RGB")

        target_size = (img.width * scale, img.height * scale)
        if upscaled_rgb.size != target_size:
            # The model upscales at its own fixed native factor - resize to
            # whatever scale the mod actually asked for.
            upscaled_rgb = upscaled_rgb.resize(target_size, Image.LANCZOS)
        upscaled_alpha = alpha.resize(target_size, Image.LANCZOS)

        result = Image.merge("RGBA", (*upscaled_rgb.split(), upscaled_alpha))
        result.save(os.path.join(out_dir, name))

    print(f"upscaled {len(names)} frame(s)")


if __name__ == "__main__":
    main()