"""Position extensions that preserve every frozen source-context lookup exactly."""

from __future__ import annotations

import torch
from torch import nn


class SegmentFactorizedPositionExtension(nn.Module):
    """Use frozen absolute rows in-range and segment/offset factors out-of-range.

    The accepted source checkpoint is an exact branch for positions below
    ``source_context``.  Only positions beyond that boundary use the new
    factorized parameters, so the extension cannot silently alter the frozen
    behavior it is intended to extend.
    """

    def __init__(
        self,
        source: nn.Embedding,
        *,
        target_context: int,
        segment_size: int,
    ) -> None:
        super().__init__()
        if target_context <= source.num_embeddings:
            raise ValueError("target context must exceed the source position table")
        if segment_size <= 0 or source.num_embeddings % segment_size:
            raise ValueError("segment size must divide the source position table")
        if target_context % segment_size:
            raise ValueError("segment size must divide target context")
        self.source_context = int(source.num_embeddings)
        self.target_context = int(target_context)
        self.segment_size = int(segment_size)
        self.embedding_dim = int(source.embedding_dim)
        self.frozen_source = nn.Embedding.from_pretrained(
            source.weight.detach().float().clone(), freeze=True,
        )

        table = self.frozen_source.weight.view(-1, segment_size, self.embedding_dim)
        local = table.mean(dim=0)
        source_age = (table - local.unsqueeze(0)).mean(dim=1)
        future_segments = target_context // segment_size - table.shape[0]
        future = torch.stack([
            source_age[index % source_age.shape[0]]
            for index in range(future_segments)
        ])
        self.local_offsets = nn.Parameter(local.clone())
        self.future_segment_age = nn.Parameter(future.clone())

    @property
    def num_embeddings(self) -> int:
        return self.target_context

    @property
    def weight(self) -> torch.Tensor:
        positions = torch.arange(self.target_context, device=self.local_offsets.device)
        return self(positions)

    def forward(self, positions: torch.Tensor) -> torch.Tensor:
        if positions.numel() and (
            int(positions.min()) < 0 or int(positions.max()) >= self.target_context
        ):
            raise IndexError("position index is outside the extended context")
        positions = positions.to(dtype=torch.long)
        source_position = positions.clamp_max(self.source_context - 1)
        source = self.frozen_source(source_position)
        future = positions >= self.source_context
        if not bool(future.any()):
            return source
        local_index = positions.remainder(self.segment_size)
        segment_index = torch.div(
            positions - self.source_context, self.segment_size, rounding_mode="floor"
        ).clamp_min(0)
        extension = self.local_offsets[local_index] + self.future_segment_age[segment_index]
        return torch.where(future.unsqueeze(-1), extension, source)

    def first_rows_exact(self, source: torch.Tensor) -> bool:
        positions = torch.arange(self.source_context, device=self.local_offsets.device)
        observed = self(positions).detach().cpu()
        return torch.equal(observed, source.detach().float().cpu())


__all__ = ["SegmentFactorizedPositionExtension"]