"""Inference for the reward-fine-tuned BERT-CNN caption reranker."""
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 CaptionReranker:
    """Pick the best of nine beam-search captions: score_k = bias[k] + weight * BERT-CNN logit."""

    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"]
        self.bias = config["rank_prior"]["bias"]
        self.weight = config["rank_prior"]["weight"]

    @torch.inference_mode()
    def logits(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")
        values = []
        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)
            values.extend(self.model(**batch).float().cpu().tolist())
        return values

    def rerank(self, visual_context, beams):
        """beams: the nine captions in beam-search order (best beam first). Returns (index, scores)."""
        if len(beams) != len(self.bias):
            raise ValueError(f"Expected exactly {len(self.bias)} captions in beam-search order")
        scores = [b + self.weight * x for b, x in zip(self.bias, self.logits(visual_context, beams))]
        best = max(range(len(scores)), key=lambda i: (scores[i], -i))  # earliest beam wins ties
        return best, scores