"""Attention reasoning on soft symbols; perception and task heads are supplied."""

from typing import NamedTuple

import torch
from torch import nn

from softasp_solver.components.evidence import EvidenceMode, clamp_evidence, restore_evidence


class ReasoningOutput(NamedTuple):
    pre_logits: torch.Tensor
    logits: torch.Tensor


def build_attention(
    d_model: int, n_heads: int, n_layers: int, dropout: float
) -> nn.TransformerEncoder:
    """Construct the existing pre-norm/GELU stack without extra initialization."""
    layer = nn.TransformerEncoderLayer(
        d_model=d_model,
        nhead=n_heads,
        dim_feedforward=d_model * 4,
        dropout=dropout,
        activation="gelu",
        batch_first=True,
        norm_first=True,
    )
    return nn.TransformerEncoder(layer, num_layers=n_layers)


def reason_from_probabilities(
    probabilities: torch.Tensor,
    *,
    projection: nn.Module,
    structure: torch.Tensor | None,
    attention: nn.Module,
    head: nn.Module,
) -> torch.Tensor:
    """Compose probability projection, optional structure, attention and task head.

    Structure can be (N,D) or (B,N,D). A head can produce per-position or pooled
    task outputs; the framework imposes no Sudoku vocabulary or board shape.
    """
    tokens = projection(probabilities)
    if structure is not None:
        tokens = tokens + (structure.unsqueeze(0) if structure.ndim == 2 else structure)
    return head(attention(tokens))


class AS2Network(nn.Module):
    """Reusable composition for an instantiation supplying its own encoder/head.

    The encoder returns (B,N,D) embeddings. Objectives are separate modules.
    The Sudoku compatibility adapter uses the same functions with its historical
    parameter ownership, avoiding checkpoint renaming or duplicate registration.
    """

    def __init__(
        self,
        *,
        encoder: nn.Module,
        concept_head: nn.Module,
        projection: nn.Module,
        attention: nn.Module,
        head: nn.Module,
        structure: nn.Module | None = None,
        training_evidence: EvidenceMode = "observed",
        inference_evidence: EvidenceMode = "predicted",
        restore_observed: bool = False,
    ) -> None:
        super().__init__()
        if inference_evidence == "observed":
            raise ValueError("Inference cannot require target symbols")
        self.encoder = encoder
        self.concept_head = concept_head
        self.projection = projection
        self.attention = attention
        self.head = head
        self.structure = structure
        self.training_evidence: EvidenceMode = training_evidence
        self.inference_evidence: EvidenceMode = inference_evidence
        self.restore_observed = restore_observed

    def forward(
        self,
        inputs: torch.Tensor,
        evidence_mask: torch.Tensor | None = None,
        observed_symbols: torch.Tensor | None = None,
        structure: torch.Tensor | None = None,
    ) -> ReasoningOutput:
        if not self.training and observed_symbols is not None:
            raise ValueError("Inference must not receive target symbols")
        pre_logits = self.concept_head(self.encoder(inputs))
        probabilities = clamp_evidence(
            pre_logits,
            evidence_mask,
            mode=self.training_evidence if self.training else self.inference_evidence,
            observed_symbols=observed_symbols,
        )
        if structure is None and self.structure is not None:
            structure = self.structure()
        logits = reason_from_probabilities(
            probabilities,
            projection=self.projection,
            structure=structure,
            attention=self.attention,
            head=self.head,
        )
        if not self.training and self.restore_observed:
            if evidence_mask is None:
                raise ValueError("Evidence restoration requires an evidence mask")
            logits = restore_evidence(logits, pre_logits, evidence_mask)
        return ReasoningOutput(pre_logits, logits)