"""Explicit text routing and learned interfaces for isolated pointer execution."""
from functools import lru_cache
import re
from typing import Any

import torch
from torch import Tensor, nn

from .number_reader import SharedNumberReader


class PromptRouter:
    """Remove the Steps line without parsing its value or changing public inputs.

    Preserve the original answer index for existing diagnostic callers. Padding
    inserted before that position is masked and RoPE positions are compacted, so
    neither count digits nor their token length reach the executor. Cached CPU
    tokenization is bounded; no prompt or activation survives in a checkpoint.
    """
    def __init__(self, tokenizer: Any) -> None:
        self.tokenizer = tokenizer
        self.executor_ids = lru_cache(maxsize=4096)(self._executor_ids)

    def _executor_ids(self, ids: tuple[int, ...]) -> tuple[int, ...]:
        prompt = self.tokenizer.decode(list(ids), skip_special_tokens=False,
                                       clean_up_tokenization_spaces=False)
        if not re.fullmatch(r"Rules: [^\n]+\nStart: [A-Z]\nSteps: [1-9][0-9]*\nAnswer:", prompt):
            raise ValueError("Isolated execution requires the raw Rules/Start/Steps/Answer prompt")
        routed = re.sub(r"\nSteps: [1-9][0-9]*", "", prompt)
        return tuple(self.tokenizer.encode(routed, add_special_tokens=False))

    def __call__(self, ids: Tensor, mask: Tensor) -> tuple[Tensor, Tensor, Tensor]:
        rows = ids.detach().cpu().tolist()
        masks = mask.detach().cpu().tolist()
        routed = torch.full_like(ids, self.tokenizer.pad_token_id)
        attention = torch.zeros_like(mask)
        for row, (tokens, valid) in enumerate(zip(rows, masks, strict=True)):
            positions = [i for i, active in enumerate(valid) if active]
            executor = self.executor_ids(tuple(tokens[i] for i in positions))
            end = positions[-1]
            if not executor or len(executor) > end + 1:
                raise ValueError("Routed executor prompt does not fit original input positions")
            routed[row, :len(executor) - 1] = torch.tensor(executor[:-1], device=ids.device)
            routed[row, end] = executor[-1]
            attention[row, :len(executor) - 1] = 1
            attention[row, end] = 1
        positions = (attention.cumsum(-1) - 1).clamp_min(0)
        return routed, attention, positions


class ReentryBridge(nn.Module):
    """Learned map from R output back to R input, applied only on re-entry.

    Normalize direction, then restore the *initial* working-state RMS scale.
    The scale contains rules/start information only and stays differentiable.
    Identity initialization of the projection does not claim T>1 equivalence.
    """
    def __init__(self, width: int) -> None:
        super().__init__()
        self.projection = nn.Linear(width, width, bias=False)
        nn.init.eye_(self.projection.weight)
        self.gain = nn.Parameter(torch.ones(width))

    def forward(self, state: Tensor, initial: Tensor) -> Tensor:
        rms = state.float().square().mean(-1, keepdim=True).add(1e-6).sqrt()
        scale = initial.float().square().mean(-1, keepdim=True).add(1e-6).sqrt()
        normalized = (state.float() / rms * scale).to(state.dtype)
        return self.projection(normalized * self.gain)


class RecurrentController(nn.Module):
    """Own learned memory; detached prompt/executor features, no numeric clock."""
    kind = 'gru'
    def __init__(self, width: int, intermediate: int = 128) -> None:
        super().__init__()
        self.intermediate = intermediate
        self.context = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, intermediate), nn.Tanh())
        self.observation = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, intermediate), nn.Tanh())
        self.cell = nn.GRUCell(intermediate, intermediate)
        self.readout = nn.Linear(intermediate, 1)

    def initialize(self, prompt_state: Tensor) -> Tensor:
        return self.context(prompt_state.detach())

    def advance(self, state: Tensor, memory: Tensor) -> tuple[Tensor, Tensor]:
        memory = self.cell(self.observation(state.detach()), memory)
        return self.readout(memory).squeeze(-1), memory


class AffineSuffixController(RecurrentController):
    """Learned prompt reading and unbounded scalar recurrence; no parsed clock.

    Initialization consumes flattened final-eight frozen token embeddings.
    The trainable cell implements a*m+b, initialized to identity (b=0), never
    a programmed decrement. Numerical labels train the cell through free-running
    recurrence; stop BCE trains its readout without distorting counting dynamics.
    Executor observations are intentionally ignored for this requested-count
    controller: timing is independent of the pointer's identity or correctness.
    """
    kind = 'affine_suffix'
    suffix_tokens = 8

    def __init__(self, width: int) -> None:
        nn.Module.__init__(self)
        self.width = width
        self.intermediate = 1
        self.context = nn.Linear(width * self.suffix_tokens, 1)
        self.observation = nn.Identity()
        self.cell = nn.Linear(1, 1)
        self.readout = nn.Linear(1, 1)
        nn.init.ones_(self.cell.weight)
        nn.init.zeros_(self.cell.bias)
        nn.init.zeros_(self.readout.weight)
        nn.init.zeros_(self.readout.bias)

    def advance(self, state: Tensor, memory: Tensor) -> tuple[Tensor, Tensor]:
        self.observation(state.detach())  # Preserve observation hooks for cache/native checks.
        memory = self.cell(memory)
        return self.readout(memory.detach()).squeeze(-1), memory


class SharedNumberController(AffineSuffixController):
    """Same learned countdown; share the number reader across positions/lengths."""
    kind = 'shared_number'

    def __init__(self, width: int) -> None:
        super().__init__(width)
        self.context = SharedNumberReader(width)


def controller_head(width: int, intermediate: int, kind: str = 'gru') -> RecurrentController:
    """Explicit architecture factory; historical checkpoints default to GRU."""
    if kind == 'gru':
        return RecurrentController(width, intermediate)
    if kind == 'affine_suffix' and intermediate == 1:
        return AffineSuffixController(width)
    if kind == 'shared_number' and intermediate == 1:
        return SharedNumberController(width)
    raise ValueError('Unknown controller kind or incompatible memory width')