"""GLiNER2.5-Decide classification graph (DESIGN.md contract), one fixed shape per instance.

Inputs:  inputs_embeds [1,N,1024] f32 (raw word-embedding rows, pad id 0 included),
         attention_mask [1,N] f32, label_routing [1,L,N] f32 (one-hot rows at the [L] markers).
Output:  logits [1,1,1,L] f32 = classifier(hidden[pos_j]); rows beyond the label count are
         classifier(0), ignored by the host.

Rewritten body = GLiNER2.5-Small's `ShapedDebertaLayer` (vendored verbatim in shaped_layer.py)
over the 24 DeBERTa-v3-large layers. `native=True` runs the unmodified transformers
`DebertaV2Model` on the same inputs, then the same routing matmul and classifier.
"""
from dataclasses import dataclass

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

from common import HIDDEN, LABEL_SLOTS
from shaped_layer import ShapedDebertaLayer, rank4_matmul

INPUT_NAMES = ["inputs_embeds", "attention_mask", "label_routing"]
OUTPUT_NAME = "logits"


@dataclass(frozen=True)
class Shape:
    seq: int
    labels: int = LABEL_SLOTS

    def contract(self):
        shapes = [(1, self.seq, HIDDEN), (1, self.seq), (1, self.labels, self.seq)]
        return dict(seq=self.seq, label_slots=self.labels,
                    inputs=[dict(name=n, shape=list(s), dtype="float32") for n, s in zip(INPUT_NAMES, shapes)],
                    output=dict(name=OUTPUT_NAME, shape=[1, 1, 1, self.labels], dtype="float32"))


def label_positions(batch):
    """[L] marker positions in request order: task order, label order, each task's [P] dropped."""
    return [p for positions in batch.schema_special_indices[0] for p in positions[1:]]


def fixed_inputs(model, captured, shape):
    batch = captured["batch"]
    n = batch.input_ids.shape[1]
    positions = label_positions(batch)
    assert n <= shape.seq and len(positions) <= shape.labels, (n, len(positions), shape)
    ids = F.pad(batch.input_ids, (0, shape.seq - n), value=model.processor.tokenizer.pad_token_id)
    attention = F.pad(batch.attention_mask, (0, shape.seq - n)).float()
    routing = torch.zeros(1, shape.labels, shape.seq)
    for j, p in enumerate(positions):
        routing[0, j, p] = 1.0
    with torch.no_grad():
        embeds = model.encoder.embeddings.word_embeddings(ids)
    return tuple(x.detach().contiguous() for x in (embeds, attention, routing))


class DecideClassifier(nn.Module):
    def __init__(self, model, shape, native=False):
        super().__init__()
        self.shape = shape
        self.native = native
        self.classifier = model.classifier
        assert [type(m).__name__ for m in self.classifier] == ["Linear", "ReLU", "Linear"]
        emb = model.encoder.embeddings
        assert emb.position_embeddings is None and emb.token_type_embeddings is None and emb.embed_proj is None
        assert model.encoder.encoder.conv is None
        if native:
            self.native_encoder = model.encoder
        else:
            self.embedding_norm = emb.LayerNorm
            with torch.no_grad():
                rel = model.encoder.encoder.get_rel_embedding().detach()
            self.layers = nn.ModuleList(ShapedDebertaLayer(x, rel, shape.seq) for x in model.encoder.encoder.layer)

    def forward(self, inputs_embeds, attention_mask, label_routing):
        if self.native:
            hidden = self.native_encoder(inputs_embeds=inputs_embeds, attention_mask=attention_mask).last_hidden_state
        else:
            hidden = self.embedding_norm(inputs_embeds) * attention_mask.unsqueeze(-1)
            mask = attention_mask[:, None, None, :] * attention_mask[:, None, :, None]
            for layer in self.layers:
                hidden = layer(hidden, mask)
        rows = rank4_matmul(label_routing.unsqueeze(1), hidden.unsqueeze(1))
        return self.classifier(rows).reshape(1, 1, 1, self.shape.labels)