#!/usr/bin/env python3
"""
Real-ESRGAN upscaling via Python/PyTorch CPU.
Usage: python3 upscale_esrgan.py <input_path> <output_path> <scale> <model>
  scale: 2 | 4
  model: 'general' | 'anime'

Strategy:
  - scale=2: single pass with RealESRGAN_x2plus (trained for 2×)
  - scale=4: two passes with RealESRGAN_x2plus (2× → 2× = 4×)
             Progressive 2×+2× produces sharper edges than single 4× pass
             because each pass extrapolates less aggressively.
  - anime:   uses RealESRGAN_x4plus_anime_6B (4× in one pass, no 2× anime model)

PyTorch CPU has no Vulkan dependency — works on Fly.dev CPU-only VMs.
"""
import sys
import os
import tempfile

def make_upsampler(models_dir, model_name, arch_blocks, model_scale):
    import torch
    from realesrgan import RealESRGANer
    from basicsr.archs.rrdbnet_arch import RRDBNet

    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64,
                    num_block=arch_blocks, num_grow_ch=32, scale=model_scale)
    model_path = os.path.join(models_dir, model_name + '.pth')

    return RealESRGANer(
        scale=model_scale,
        model_path=model_path,
        model=model,
        tile=512,
        tile_pad=32,
        pre_pad=0,
        half=False,   # FP16 only on GPU
        device=torch.device('cpu'),
    )

def enhance(upsampler, img_np):
    """Run one pass. img_np is RGB uint8 numpy array. Returns RGB uint8."""
    import numpy as np
    bgr = img_np[:, :, ::-1]
    out_bgr, _ = upsampler.enhance(bgr, outscale=upsampler.scale)
    return out_bgr[:, :, ::-1]

def main():
    if len(sys.argv) != 5:
        print(f"Usage: {sys.argv[0]} <input> <output> <scale> <model>", file=sys.stderr)
        sys.exit(1)

    input_path  = sys.argv[1]
    output_path = sys.argv[2]
    scale       = int(sys.argv[3])
    model_type  = sys.argv[4]  # 'general' or 'anime'

    import numpy as np
    from PIL import Image

    models_dir = os.environ.get(
        'ESRGAN_PYTHON_MODELS_PATH',
        os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'models')
    )

    img = Image.open(input_path).convert('RGB')
    img_np = np.array(img)

    if model_type == 'anime':
        # No x2plus anime model — use x4plus_anime_6B for both 2× and 4×
        upsampler = make_upsampler(models_dir, 'RealESRGAN_x4plus_anime_6B',
                                   arch_blocks=6, model_scale=4)
        if scale == 2:
            # Run 4× then downscale to 2× (best quality for anime 2×)
            out = enhance(upsampler, img_np)
            h, w = img_np.shape[:2]
            out = np.array(Image.fromarray(out).resize((w * 2, h * 2), Image.LANCZOS))
        else:
            out = enhance(upsampler, img_np)

    else:
        # General: use x2plus model
        # scale=2 → one pass; scale=4 → two passes (progressive upscaling)
        upsampler2x = make_upsampler(models_dir, 'RealESRGAN_x2plus',
                                     arch_blocks=23, model_scale=2)
        out = enhance(upsampler2x, img_np)          # first 2× pass
        if scale == 4:
            out = enhance(upsampler2x, out)          # second 2× pass → 4× total

    Image.fromarray(out).save(output_path)
    print("OK", flush=True)

if __name__ == '__main__':
    main()