#!/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.

Resolves <model-name> to $UPSCALE_MODELS_DIR/<model-name>.onnx and picks
single-pass or tiled inference from that model's own declared ONNX input
shape - a model with a fixed height/width (e.g. 128x128) needs tiling for
any image bigger than that; one with symbolic dimensions runs the whole
image through in one pass, whatever its size.

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

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


def extend_opaque_color(rgb, alpha, threshold=128):
    """Replaces every below-threshold-alpha pixel's RGB with its nearest
    opaque neighbor's, so the model (RGB-only, see run_tile()) sees a
    smooth color field instead of a hard edge against flat black -
    otherwise its convolutions blend genuine edge colors with that black
    background, leaving dark fringing on the final (correctly re-masked)
    output."""
    transparent = np.asarray(alpha) < threshold
    if not transparent.any() or transparent.all():
        return rgb
    _, indices = ndimage.distance_transform_edt(transparent, return_indices=True)
    arr = np.asarray(rgb)
    filled = arr[indices[0], indices[1]]
    return Image.fromarray(np.where(transparent[..., None], filled, arr))


def run_tile(session, input_name, tile_rgb):
    """Runs one tile (already at the model's required size) and returns its
    upscaled RGB as a uint8 HxWx3 array - handles both the with-batch-dim
    ([1,3,H,W]) and without ([3,H,W]) output conventions different ONNX
    exports use."""
    arr = tile_rgb.astype(np.float32) / 255.0
    arr = np.transpose(arr, (2, 0, 1))[np.newaxis, ...]
    out = session.run(None, {input_name: arr})[0]
    if out.ndim == 4:
        out = out[0]
    out = np.clip(out, 0.0, 1.0) * 255.0
    return np.transpose(out, (1, 2, 0)).astype(np.uint8)


def upscale_dynamic(session, input_name, rgb_image):
    return Image.fromarray(run_tile(session, input_name, np.asarray(rgb_image)))


def upscale_tiled(session, input_name, rgb_image, tile_size, native_scale):
    """Splits an image larger than the model's fixed input size into
    overlapping tile_size x tile_size tiles (edge-padding any tile that
    would otherwise run past the image bounds), runs each individually, and
    stitches the results back into one canvas - later tiles simply
    overwrite the small overlap region of earlier ones, which is invisible
    at sprite scale."""
    src = np.asarray(rgb_image)
    h, w = src.shape[:2]
    canvas = np.zeros((h * native_scale, w * native_scale, 3), dtype=np.uint8)
    overlap = 16
    stride = tile_size - overlap

    y = 0
    while True:
        ty = min(y, max(0, h - tile_size))
        x = 0
        while True:
            tx = min(x, max(0, w - tile_size))
            tile = src[ty:ty + tile_size, tx:tx + tile_size]
            th, tw = tile.shape[:2]
            if th < tile_size or tw < tile_size:
                tile = np.pad(tile, ((0, tile_size - th), (0, tile_size - tw), (0, 0)), mode="edge")
            out = run_tile(session, input_name, tile)[: th * native_scale, : tw * native_scale]
            oy, ox = ty * native_scale, tx * native_scale
            canvas[oy:oy + th * native_scale, ox:ox + tw * native_scale] = out
            if tx + tile_size >= w:
                break
            x += stride
        if ty + tile_size >= h:
            break
        y += stride
    return Image.fromarray(canvas)


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

    models_dir = os.environ["UPSCALE_MODELS_DIR"]
    model_path = os.path.join(models_dir, f"{model_name}.onnx")
    if not os.path.isfile(model_path):
        available = sorted(f[:-5] for f in os.listdir(models_dir) if f.endswith(".onnx"))
        sys.exit(f"unknown upscale model {model_name!r} - available: {', '.join(available)}")

    session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
    input_info = session.get_inputs()[0]
    input_name = input_info.name
    # A fixed (integer) height/width means this model only ever accepts
    # that exact input size - a dynamic one (reported as a symbolic name
    # such as "height") accepts any size in one pass. For a fixed-size
    # model, its own declared output shape gives the native scale factor
    # directly, with no need to run it just to find out.
    fixed_size = input_info.shape[2] if isinstance(input_info.shape[2], int) else None
    native_scale = session.get_outputs()[0].shape[-1] // fixed_size if fixed_size else None

    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]
        rgb = extend_opaque_color(rgb, alpha)

        # 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.
        if fixed_size is not None and (img.width > fixed_size or img.height > fixed_size):
            upscaled_rgb = upscale_tiled(session, input_name, rgb, fixed_size, native_scale)
        elif fixed_size is not None:
            padded = np.pad(np.asarray(rgb), ((0, fixed_size - img.height), (0, fixed_size - img.width), (0, 0)), mode="edge")
            out = run_tile(session, input_name, padded)
            upscaled_rgb = Image.fromarray(out[: img.height * native_scale, : img.width * native_scale])
        else:
            upscaled_rgb = upscale_dynamic(session, input_name, 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()