"""ShadeNet-3.2 torch inference (standalone CLI).

Usage:
    pip install -r requirements.txt
    python inference.py input.jpg --output_dir ./output
    python inference.py ./photos --checkpoint ./checkpoints/shadenet32.ckpt
"""
import argparse
import os

import torch
from PIL import Image

from model import load_shadenet32
from inference_utils import build_grid, pil_to_np, resize_pad

IMAGE_SIZE = 512
ENS_SCALES = (0.875, 1.0, 1.125)


@torch.no_grad()
def run(model, img_rgb: Image.Image, device: str,
        image_size: int = IMAGE_SIZE, scales=ENS_SCALES) -> "torch.Tensor":
    """3-pass multi-scale median (matches the model card / Space)."""
    outs = []
    for s in scales:
        side = max(16, int(round(image_size * s)) // 16 * 16)  # dict needs /16
        scaled = resize_pad(img_rgb, side)
        x = torch.from_numpy(pil_to_np(scaled)).to(device)
        o = model(x)
        if side != image_size:
            o = torch.nn.functional.interpolate(
                o, size=(image_size, image_size), mode="bilinear",
                align_corners=False)
        outs.append(o.cpu())
    return torch.stack(outs, dim=0).median(dim=0).values


def main():
    ap = argparse.ArgumentParser(description="ShadeNet-3.2 inference")
    ap.add_argument("input", help="Image file or folder")
    ap.add_argument("--checkpoint", default="./checkpoints/shadenet32.ckpt")
    ap.add_argument("--output_dir", default="./output")
    ap.add_argument("--image-size", type=int, default=IMAGE_SIZE)
    ap.add_argument("--no-ema", action="store_true",
                    help="Use raw weights instead of EMA")
    args = ap.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = load_shadenet32(args.checkpoint, device, use_ema=not args.no_ema)

    if os.path.isdir(args.input):
        files = sorted(os.path.join(args.input, f) for f in os.listdir(args.input)
                       if f.lower().endswith((".png", ".jpg", ".jpeg", ".webp")))
    else:
        files = [args.input]
    os.makedirs(args.output_dir, exist_ok=True)
    for fp in files:
        img_rgb = resize_pad(Image.open(fp), args.image_size)
        out = run(model, img_rgb, device).numpy()
        base = os.path.splitext(os.path.basename(fp))[0]
        build_grid(img_rgb, out).save(
            os.path.join(args.output_dir, f"{base}_result.png"))
        print(f"Saved: {base}_result.png")


if __name__ == "__main__":
    main()