"""Inference for exported BERT-CNN and ModernBERT-CNN relatedness models."""
import json
from pathlib import Path

import torch
from safetensors.torch import load_file
from transformers import AutoTokenizer

from modeling import VisualRelatednessModel, backbone_config_from_dict


class RelatednessScorer:
    def __init__(self, model_path, device="cpu"):
        path = Path(model_path)
        if not path.is_dir():
            from huggingface_hub import snapshot_download
            path = Path(snapshot_download(str(model_path), allow_patterns=["config.json", "model.safetensors", "tokenizer/*"]))
        config = json.loads((path / "config.json").read_text())
        self.device = torch.device(device)
        self.model = VisualRelatednessModel(backbone_config_from_dict(config["bert_config"]), **config["model_options"])
        self.model.load_state_dict(load_file(str(path / "model.safetensors")))
        self.model.to(self.device).eval()
        self.tokenizer = AutoTokenizer.from_pretrained(path / "tokenizer", local_files_only=True, use_fast=True)
        self.tokenizer.padding_side = "right"
        self.seq_len = config["model_options"]["seq_len"]

    @torch.inference_mode()
    def score(self, visual_context, captions, batch_size=16):
        if not isinstance(visual_context, str) or not visual_context.strip():
            raise ValueError("Provide image-context words, such as dog, ball, grass")
        if not captions or any(not isinstance(c, str) or not c.strip() for c in captions):
            raise ValueError("Provide at least one nonempty caption")
        if batch_size < 1:
            raise ValueError("batch_size must be positive")
        scores = []
        for start in range(0, len(captions), batch_size):
            selected = captions[start:start + batch_size]
            batch = self.tokenizer([visual_context] * len(selected), selected, padding="max_length",
                                   truncation="longest_first", max_length=self.seq_len,
                                   return_tensors="pt").to(self.device)
            scores.extend(self.model.probabilities(self.model(**batch)).float().cpu().tolist())
        if not all(0 <= score <= 1 for score in scores):
            raise RuntimeError("Model produced invalid scores")
        return scores