"""Two-page spread vs. single page classifier for document scans.

    python inference.py scan1.jpg scan2.jpg ...

Prints the logit and the decision for each scan.
"""
import json
import sys
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
from PIL import Image, ImageOps
from safetensors.torch import load_file
from torchvision.models import efficientnet_b0

HERE = Path(__file__).resolve().parent


class SpreadNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = efficientnet_b0(weights=None).features
        self.cls = nn.Sequential(nn.Dropout(0.3), nn.Linear(1280, 1))

    def forward(self, x):
        return self.cls(self.features(x).mean((2, 3))).squeeze(1)


def load_model(folder=HERE):
    cfg = json.loads((Path(folder) / "config.json").read_text())
    model = SpreadNet()
    model.load_state_dict(load_file(str(Path(folder) / "model.safetensors")))
    return model.eval(), cfg


def preprocess(path, cfg):
    """RGB, longest side = image_size, centred on a gray square, ImageNet normalisation."""
    size = cfg["image_size"]
    im = ImageOps.exif_transpose(Image.open(path)).convert("RGB")
    k = size / max(im.size)
    im = im.resize((max(1, round(im.width * k)), max(1, round(im.height * k))), Image.LANCZOS)
    sq = Image.new("RGB", (size, size), (cfg["pad_value"],) * 3)
    sq.paste(im, ((size - im.width) // 2, (size - im.height) // 2))
    x = np.asarray(sq, np.float32) / 255
    x = (x - np.array(cfg["mean"], np.float32)) / np.array(cfg["std"], np.float32)
    return torch.from_numpy(x.transpose(2, 0, 1))


@torch.inference_mode()
def predict(paths, model, cfg):
    """Logits: above recommended_logit_threshold means a two-page spread."""
    return model(torch.stack([preprocess(p, cfg) for p in paths])).tolist()


if __name__ == "__main__":
    model, cfg = load_model()
    thr = cfg["recommended_logit_threshold"]
    for path, z in zip(sys.argv[1:], predict(sys.argv[1:], model, cfg)):
        print(f"{path}\t{z:.2f}\t{'spread' if z > thr else 'single'}")