"""Official interface for the compact auto-routed U5d candidate."""

from __future__ import annotations

from pathlib import Path
from typing import Sequence

import torch

from modchallenge.interface.base_model import ModularMultiplicationModel
from src.adaptive_algorithm_router_u3 import RawInputAlgorithmRouter
from src.compact_auto_routed_short_u5d import (
    CompactAutoRoutedShortU5d,
    CompactSparseUniversalShortU5d,
)


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 CompactAutoRoutedShortU5dSubmission(ModularMultiplicationModel):
    """Learned raw routing into compact Montgomery and short-Horner programs."""

    def __init__(self) -> None:
        self.model: CompactAutoRoutedShortU5d | 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,
        )
        executor = CompactSparseUniversalShortU5d(
            **payload["model_config"]["executor"]
        ).to(self.device)
        router = RawInputAlgorithmRouter(
            **payload["model_config"]["router"]
        ).to(self.device)
        self.model = CompactAutoRoutedShortU5d(
            executor=executor, router=router
        ).to(self.device)
        result = self.model.load_state_dict(payload["state_dict"], strict=True)
        if result.missing_keys or result.unexpected_keys:
            raise RuntimeError("compact U5d 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__ = ["CompactAutoRoutedShortU5dSubmission"]