"""Command-line inference for the Ermolaev historical print style model."""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from PIL import Image, ImageOps

from model import load_generator


REPO_ID = "tuva/Ermolaev"
DEFAULT_WEIGHTS = "weights/ermolaev_epoch2_recommended.pth"


def stylize(source, destination, weights=DEFAULT_WEIGHTS, max_side=1600, device=None):
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    weights_path = hf_hub_download(REPO_ID, weights)
    generator = load_generator(weights_path, device)

    with Image.open(source) as raw:
        image = ImageOps.exif_transpose(raw).convert("RGB")
    scale = min(1.0, max_side / max(image.size))
    if scale < 1:
        image = image.resize(tuple(round(v * scale) for v in image.size), Image.Resampling.LANCZOS)
    array = np.asarray(image, dtype=np.float32).copy() / 127.5 - 1.0
    tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0).to(device)
    height, width = tensor.shape[-2:]
    pad_h, pad_w = (-height) % 4, (-width) % 4
    if pad_h or pad_w:
        tensor = F.pad(tensor, (0, pad_w, 0, pad_h), mode="reflect")
    with torch.inference_mode():
        output = generator(tensor)[0, :, :height, :width]
    output = ((output.clamp(-1, 1) + 1) * 127.5).permute(1, 2, 0).byte().cpu().numpy()
    result = Image.fromarray(output)
    destination = Path(destination)
    destination.parent.mkdir(parents=True, exist_ok=True)
    result.save(destination)
    return destination


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("input")
    parser.add_argument("output")
    parser.add_argument("--weights", default=DEFAULT_WEIGHTS)
    parser.add_argument("--max-side", type=int, default=1600)
    args = parser.parse_args()
    print(stylize(args.input, args.output, args.weights, args.max_side))