"""Official interface for compiled-local Montgomery v42."""

from __future__ import annotations

from pathlib import Path
from typing import Sequence

import torch

from modchallenge.interface.base_model import ModularMultiplicationModel
from src.compact_sparse_montgomery_v39 import CompiledLocalSparseMontgomeryV42


MAXIMUM_WIDTH = 512


def _hex_digits(value: str) -> list[int]:
    number = int(value)
    if number < 0:
        raise ValueError("challenge arguments must be non-negative")
    width = max(1, (number.bit_length() + 3) // 4)
    return [
        (number >> (4 * position)) & 15
        for position in range(width - 1, -1, -1)
    ]


class CompiledLocalMontgomeryV42Submission(ModularMultiplicationModel):
    """Raw arguments through v20 arithmetic and compiled learned local maps."""

    def __init__(self) -> None:
        self.model: CompiledLocalSparseMontgomeryV42 | None = None
        self.device: torch.device | None = None

    def load(self, model_dir: str) -> None:
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        if self.device.type == "cpu":
            torch.set_num_threads(4)
        payload = torch.load(
            Path(model_dir) / "weights.pt",
            map_location=self.device,
            weights_only=True,
        )
        self.model = CompiledLocalSparseMontgomeryV42(
            **payload["model_config"]
        ).to(self.device)
        result = self.model.load_state_dict(payload["state_dict"], strict=True)
        if result.missing_keys or result.unexpected_keys:
            raise RuntimeError("v42 checkpoint did not load strictly")
        self.model.eval()

    def preprocess_a(self, a: str) -> list[int]:
        return _hex_digits(a)

    def preprocess_b(self, b: str) -> list[int]:
        return _hex_digits(b)

    def preprocess_p(self, p: str) -> list[int]:
        return _hex_digits(p)

    @torch.inference_mode()
    def predict_digits(self, a_enc, b_enc, p_enc) -> list[int]:
        return self.predict_digits_batch([(a_enc, b_enc, p_enc)])[0]

    @torch.inference_mode()
    def predict_digits_batch(self, inputs: Sequence[tuple]) -> list[list[int]]:
        if self.model is None:
            raise RuntimeError("model has not been loaded")
        answers: list[list[int] | None] = [None] * len(inputs)
        supported = [
            index
            for index, (_a, _b, modulus) in enumerate(inputs)
            if len(modulus) <= MAXIMUM_WIDTH
        ]
        if supported:
            prediction = self.model.predict_encoded(
                [inputs[index][0] for index in supported],
                [inputs[index][1] for index in supported],
                [inputs[index][2] for index in supported],
            )
            for index, row in zip(supported, prediction.answers.tolist()):
                answers[index] = [int(digit) for digit in row]
        for index, answer in enumerate(answers):
            if answer is None:
                answers[index] = [0]
        return [answer for answer in answers if answer is not None]

    def max_batch_size(self) -> int:
        return 20


__all__ = ["CompiledLocalMontgomeryV42Submission"]