"""The learned reduction cell: a 2-layer GRU that approximates

    cell(state, m, digit) = (state * R + m * digit) mod p

entirely through trained parameters. No hand-coded modulus reduction of the
challenge prime appears anywhere in this file -- the reduction is produced
solely by embeddings + GRU + linear head + softmax/argmax.

State, m, and p are all explicit digit/bit representations (not opaque
latent vectors) fed in and read out as integers via `mac_cell.digits`.
"""

from __future__ import annotations

import torch
import torch.nn as nn

# Historical default matching task 2's T1/T2-only domain (p < 2**8).
# ReductionCell now takes p_bits as a constructor argument instead (task 3
# extends the domain to p up to 2**16 for T3); this constant is kept only
# as that argument's default so old call sites without an explicit p_bits
# keep behaving exactly as before.
P_BITS = 8


class ReductionCell(nn.Module):
    """Two-layer unidirectional GRU operating on explicit digit-state input.

    Per-step input layout (last dim, size `input_size`):
        [digit_in_embed | state_digits_embed (D slots) | m_digits_embed (D slots) | p_bits | phase_flag (1)]

    Output per step: logits of shape (..., state_digits, radix) predicting
    the *next* state's digits.
    """

    def __init__(
        self,
        radix: int,
        state_digits: int,
        hidden_size: int,
        digit_embed_dim: int = 16,
        num_layers: int = 2,
        p_bits: int = P_BITS,
    ) -> None:
        super().__init__()
        self.radix = radix
        self.state_digits = state_digits
        self.hidden_size = hidden_size
        self.digit_embed_dim = digit_embed_dim
        self.num_layers = num_layers
        self.p_bits = p_bits

        # Shared embedding table for any base-R digit value (input stream
        # digits, state digits, m digits alike -- they are all "a digit in
        # [0, R)" semantically).
        self.digit_embed = nn.Embedding(radix, digit_embed_dim)

        self.input_size = (
            digit_embed_dim  # digit_in
            + state_digits * digit_embed_dim  # state_digits
            + state_digits * digit_embed_dim  # m_digits
            + p_bits  # p bits, raw 0/1 floats
            + 1  # phase flag, raw 0.0/1.0 float
        )

        self.gru = nn.GRU(
            input_size=self.input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
        )
        self.head = nn.Linear(hidden_size, state_digits * radix)

    def num_params(self) -> int:
        return sum(p.numel() for p in self.parameters())

    def assemble_features(
        self,
        digit_in: torch.Tensor,      # (..., ) long
        state_digits: torch.Tensor,  # (..., D) long
        m_digits: torch.Tensor,      # (..., D) long
        p_bits: torch.Tensor,        # (..., 8) float
        phase_flag: torch.Tensor,    # (..., 1) float
    ) -> torch.Tensor:
        """Build the per-step input feature vector. Works for a full
        (batch, T, ...) sequence or a single step (batch, ...) alike -- the
        embedding + concatenation ops broadcast over any leading dims.
        """
        digit_in_e = self.digit_embed(digit_in)  # (..., E)
        state_e = self.digit_embed(state_digits)  # (..., D, E)
        state_e = state_e.flatten(start_dim=-2)  # (..., D*E)
        m_e = self.digit_embed(m_digits)
        m_e = m_e.flatten(start_dim=-2)  # (..., D*E)
        return torch.cat([digit_in_e, state_e, m_e, p_bits, phase_flag], dim=-1)

    def forward_sequence(
        self, features: torch.Tensor, h0: torch.Tensor | None = None
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Run the GRU over a full (batch, T, input_size) sequence at once.

        Used for teacher-forced training, where every step's true
        state_digits/m_digits are already known and pre-assembled into
        `features`.

        Returns (logits (batch, T, D, R), final hidden state).
        """
        out, h_n = self.gru(features, h0)
        logits = self.head(out)
        batch, seq_len, _ = logits.shape
        logits = logits.view(batch, seq_len, self.state_digits, self.radix)
        return logits, h_n

    def forward_step(
        self, features: torch.Tensor, h: torch.Tensor | None
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Run a single step: `features` is (batch, input_size).

        Used for autoregressive rollout at inference time, where each
        step's state_digits input is the model's own previous prediction.

        Returns (logits (batch, D, R), new hidden state).
        """
        out, h_n = self.gru(features.unsqueeze(1), h)
        logits = self.head(out.squeeze(1))
        logits = logits.view(-1, self.state_digits, self.radix)
        return logits, h_n