"""ShadeNet-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/shadenet2.ckpt
"""
import argparse
import os

import torch
from PIL import Image

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

IMAGE_SIZE = 512


@torch.no_grad()
def run(model, img_rgb: Image.Image, device: str) -> "torch.Tensor":
    x = torch.from_numpy(pil_to_np(img_rgb)).to(device)
    return model(x).detach().cpu()


def main():
    ap = argparse.ArgumentParser(description="ShadeNet-2 inference")
    ap.add_argument("input", help="Image file or folder")
    ap.add_argument("--checkpoint", default="./checkpoints/shadenet2.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_shadenet2(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()