"""Convert castorini/monot5-small-msmarco-10k into an ONNX graph that survives Teradata's
ONNXSeq2Seq function, which only ever returns tokenizer-decoded text and has no native
concept of a numeric score.

The trick: add 1000 dedicated vocabulary tokens for quantized score buckets
(<bucket_000> .. <bucket_999>), compute the real relevance score via one genuine T5
forward pass, and output the id of the matching bucket token directly -- no beam search
needed, no dependency on any third-party conversion package. See README.md for the full
mechanism and why it works against ONNXSeq2Seq's actual (undocumented) behavior.
"""

from pathlib import Path

import torch
import torch.nn as nn
from transformers import AutoTokenizer, T5ForConditionalGeneration

UPSTREAM_REPO = "castorini/monot5-small-msmarco-10k"
N_BUCKETS = 1000
OPSET = 17


class MonoT5Score(nn.Module):
    """Wraps a real T5 encoder + one real decoder step; outputs the id of a dedicated
    vocabulary token whose decoded text is the quantized relevance score, e.g. '<bucket_042>'.
    """

    def __init__(self, model: T5ForConditionalGeneration, true_id: int, false_id: int, bucket_ids: list[int]):
        super().__init__()
        self.model = model
        self.true_id = true_id
        self.false_id = false_id
        self.register_buffer("bucket_token_ids", torch.tensor(bucket_ids, dtype=torch.int64))
        self.n_buckets = len(bucket_ids)

    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        batch = input_ids.shape[0]
        decoder_input_ids = torch.zeros((batch, 1), dtype=torch.long, device=input_ids.device)
        out = self.model(input_ids=input_ids, attention_mask=attention_mask, decoder_input_ids=decoder_input_ids)
        logits = out.logits[:, 0, :]
        two = torch.stack([logits[:, self.false_id], logits[:, self.true_id]], dim=1)
        probs = torch.softmax(two, dim=1)
        p_true = probs[:, 1]
        bucket_idx = torch.clamp((p_true * (self.n_buckets - 1)).round().long(), 0, self.n_buckets - 1)
        token_ids = self.bucket_token_ids[bucket_idx]
        return token_ids.view(batch, 1)


def bucket_token_name(index: int, n_buckets: int) -> str:
    width = len(str(n_buckets - 1))
    return f"<bucket_{index:0{width}d}>"


def main() -> None:
    repo_dir = Path(__file__).parent
    onnx_dir = repo_dir / "onnx"
    onnx_dir.mkdir(exist_ok=True)

    tok = AutoTokenizer.from_pretrained(UPSTREAM_REPO)
    bucket_tokens = [bucket_token_name(i, N_BUCKETS) for i in range(N_BUCKETS)]
    tok.add_tokens(bucket_tokens)
    tok.save_pretrained(repo_dir)
    tok = AutoTokenizer.from_pretrained(repo_dir)  # reload so the new bucket tokens are active

    model = T5ForConditionalGeneration.from_pretrained(UPSTREAM_REPO)
    model.eval()
    model.config.save_pretrained(repo_dir)
    model.generation_config.save_pretrained(repo_dir)

    true_id = tok.convert_tokens_to_ids("▁true")
    false_id = tok.convert_tokens_to_ids("▁false")
    bucket_ids = [tok.convert_tokens_to_ids(bucket_token_name(i, N_BUCKETS)) for i in range(N_BUCKETS)]

    score_model = MonoT5Score(model, true_id, false_id, bucket_ids)
    score_model.eval()

    sample = tok("Query: hello world Document: foo bar baz Relevant:", return_tensors="pt")
    onnx_path = onnx_dir / "model.onnx"
    torch.onnx.export(
        score_model,
        (sample.input_ids, sample.attention_mask),
        str(onnx_path),
        input_names=["input_ids", "attention_mask"],
        output_names=["relevance_token_id"],
        dynamic_axes={
            "input_ids": {0: "batch", 1: "sequence"},
            "attention_mask": {0: "batch", 1: "sequence"},
            "relevance_token_id": {0: "batch"},
        },
        opset_version=OPSET,
        dynamo=False,  # the newer dynamo-based exporter needs the optional onnxscript
        # package; the legacy TorchScript-based exporter doesn't have that dependency.
    )
    print(f"wrote {onnx_path} (opset {OPSET}, {N_BUCKETS} bucket tokens added to vocab)")


if __name__ == "__main__":
    main()