"""V20-sized learned controller for the two-program U2 arithmetic board."""

from __future__ import annotations

from collections.abc import Sequence

import torch
import torch.nn as nn
import torch.nn.functional as F

from src.self_clocked_parallel_montgomery_teacher_student import select_candidate
from src.universal_self_clocked_u1 import OP_COUNT, PHASE_COUNT, UniversalU1Memory
from src.universal_self_clocked_u2 import (
    ALGORITHM_COUNT,
    UniversalSelfClockedU2,
)


class CompactMixedAlgorithmSelfClockTeacher(nn.Module):
    """Learn a finite program using compact scalar clock transitions."""

    def __init__(self, *, maximum_width: int = 512, maximum_clock: int = 2048) -> None:
        super().__init__()
        if maximum_width < 1:
            raise ValueError("maximum_width must be positive")
        if maximum_clock < 4 * maximum_width:
            raise ValueError("maximum_clock must cover four bits per base-16 digit")
        self.maximum_width = int(maximum_width)
        self.maximum_clock = int(maximum_clock)
        self.operation = nn.Embedding(PHASE_COUNT * 2, OP_COUNT)
        self.initial_phase = nn.Embedding(ALGORITHM_COUNT, PHASE_COUNT)
        self.initial_clock_affine = nn.Embedding(ALGORITHM_COUNT, 2)
        self.decrement_clock_affine = nn.Linear(1, 1)
        self.terminal_clock = nn.Parameter(torch.empty(1))
        self.terminal_test = nn.Linear(1, 2)
        self.next_phase = nn.Embedding(PHASE_COUNT * OP_COUNT, PHASE_COUNT)
        for parameter in self.parameters():
            nn.init.normal_(parameter, std=0.02)
        self.config = {
            "maximum_width": self.maximum_width,
            "maximum_clock": self.maximum_clock,
        }

    @property
    def parameter_count(self) -> int:
        return sum(parameter.numel() for parameter in self.parameters())

    def _validate_algorithms(self, algorithm_ids: torch.Tensor) -> None:
        if algorithm_ids.ndim != 1:
            raise ValueError("algorithm_ids must be a one-dimensional batch")
        if bool(
            ((algorithm_ids < 0) | (algorithm_ids >= ALGORITHM_COUNT)).any().item()
        ):
            raise ValueError("algorithm id is outside the learned instruction domain")

    def initial_clock_values(
        self, algorithm_ids: torch.Tensor, widths: torch.Tensor
    ) -> torch.Tensor:
        self._validate_algorithms(algorithm_ids)
        if algorithm_ids.shape != widths.shape:
            raise ValueError("algorithm ids and widths must align")
        scale_bias = self.initial_clock_affine(algorithm_ids)
        return scale_bias[:, 0] * widths.to(scale_bias.dtype) + scale_bias[:, 1]

    def raw_decrement_clock_values(self, clocks: torch.Tensor) -> torch.Tensor:
        return self.decrement_clock_affine(
            clocks.to(torch.float32).unsqueeze(-1)
        ).squeeze(-1)

    def decrement_clock_values(self, clocks: torch.Tensor) -> torch.Tensor:
        return F.relu(self.raw_decrement_clock_values(clocks))

    def terminal_test_logits(self, clocks: torch.Tensor) -> torch.Tensor:
        return self.terminal_test(clocks.to(torch.float32).unsqueeze(-1))

    def _discrete_clock(self, values: torch.Tensor) -> torch.Tensor:
        return values.round().clamp(0, self.maximum_clock).to(torch.long)

    def start(
        self, width: int, algorithm_ids: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if not 1 <= int(width) <= self.maximum_width:
            raise ValueError("width exceeds the learned controller domain")
        self._validate_algorithms(algorithm_ids)
        widths = torch.full_like(algorithm_ids, int(width))
        phase = self.initial_phase(algorithm_ids).argmax(-1)
        clock = self._discrete_clock(
            self.initial_clock_values(algorithm_ids, widths)
        )
        return phase, clock

    def operation_logits(
        self, phases: torch.Tensor, clocks: torch.Tensor
    ) -> torch.Tensor:
        clock_class = self.terminal_test_logits(clocks).argmax(-1)
        return self.operation(phases * 2 + clock_class)

    def instruction(self, phase: torch.Tensor, clock: torch.Tensor) -> torch.Tensor:
        return self.operation_logits(phase, clock).argmax(-1)

    def advance(
        self,
        phase: torch.Tensor,
        clock: torch.Tensor,
        operation: torch.Tensor,
        algorithm_ids: torch.Tensor,
        *,
        width: int,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        widths = torch.full_like(clock, int(width))
        row_clock = self._discrete_clock(self.decrement_clock_values(clock))
        reset_clock = self._discrete_clock(
            self.initial_clock_values(algorithm_ids, widths)
        )
        terminal_clock = self._discrete_clock(self.terminal_clock).expand_as(clock)
        next_clock = select_candidate(
            (
                row_clock,
                reset_clock,
                terminal_clock,
                row_clock,
                terminal_clock,
                clock,
            ),
            operation,
        )
        next_phase = self.next_phase(phase * OP_COUNT + operation).argmax(-1)
        return next_phase, next_clock


class V20CompactUniversalSelfClockedU2(UniversalSelfClockedU2):
    """Use the U2 arithmetic board with a compact learned controller."""

    def __init__(
        self,
        *,
        reducer: dict,
        controller: dict,
        mac: dict,
        subtract: dict,
        factor: dict,
        cancellation_teacher: dict,
        prime_two_router: dict,
    ) -> None:
        super().__init__(
            reducer=reducer,
            controller=controller,
            mac=mac,
            subtract=subtract,
            factor=factor,
            cancellation_teacher=cancellation_teacher,
            prime_two_router=prime_two_router,
        )
        self.controller = CompactMixedAlgorithmSelfClockTeacher(**controller)
        self.config = {
            "reducer": dict(reducer),
            "controller": dict(controller),
            "mac": dict(mac),
            "subtract": dict(subtract),
            "factor": dict(factor),
            "cancellation_teacher": dict(cancellation_teacher),
            "prime_two_router": dict(prime_two_router),
        }

    def _controller_instruction(self, memory: UniversalU1Memory) -> torch.Tensor:
        return self.controller.instruction(memory.phase, memory.clock)


__all__ = [
    "CompactMixedAlgorithmSelfClockTeacher",
    "V20CompactUniversalSelfClockedU2",
]