"""Development-only frozen parent with learned service/log/base head selection."""
import torch
from torch import nn

from ._window_gate import WindowGateModel


class SourceKindGateModel(WindowGateModel):
    """Same parent forward/packing; two CLS gates choose one head per window."""

    def __init__(self, encoder):
        super().__init__(encoder)
        self.log_head = nn.Linear(encoder.config.hidden_size, 1, bias=False)
        self.log_gate = nn.Linear(encoder.config.hidden_size, 1)

    def choose_head(self, pooled, cls):
        gate = self.window_gate(cls)
        original = self.head(pooled)
        service = self.service_head(pooled)
        log_gate = self.log_gate(cls)
        log = self.log_head(pooled) + self.head.bias
        valid = (torch.isfinite(gate) & torch.isfinite(log_gate) & torch.isfinite(original)
                 & torch.isfinite(service) & torch.isfinite(log))
        chosen = torch.where(gate >= 0, service, torch.where(log_gate >= 0, log, original))
        return torch.where(valid, chosen, torch.full_like(chosen, float('nan')))


def self_check():
    """Synthetic CPU branch/pooling/state regression; no trained encoder or CUDA."""
    from types import SimpleNamespace

    class Encoder(nn.Module):
        config = SimpleNamespace(hidden_size=1)

        def forward(self, input_ids, **kwargs):
            return SimpleNamespace(last_hidden_state=input_ids.float().unsqueeze(-1))

    model = SourceKindGateModel(Encoder()).eval()
    with torch.no_grad():
        model.head.weight.fill_(10); model.head.bias.fill_(2)
        model.service_head.weight.fill_(-10); model.service_head.bias.fill_(3)
        model.window_gate.weight.fill_(1); model.window_gate.bias.zero_()
        model.log_head.weight.fill_(4)
        model.log_gate.weight.fill_(-1); model.log_gate.bias.fill_(-1)
    for first, expected in [(-2., 4.), (-1., 4.), (-.5, 7.), (0., -2.), (2., -2.)]:
        inputs = torch.tensor([[first, -1., 2.]])
        result = model(inputs, torch.ones_like(inputs), [(1, 3)])
        assert result.tolist() == [expected]
        states = inputs.unsqueeze(-1)
        tokens = model.choose_head(states, states[:, :1, :]).squeeze(-1)
        assert torch.equal(tokens[:, 1:3].mean(1), result)
    state = {k: v.detach().clone() for k, v in model.state_dict().items()}
    for key in state:
        model.load_state_dict(state)
        with torch.no_grad():
            dict(model.named_parameters())[key].fill_(float('inf'))
        try:
            model(torch.tensor([[-.5, -1., 2.]]), torch.ones(1, 3), [(1, 3)])
        except ValueError as error:
            assert str(error) == 'nonfinite_scores'
        else:
            raise AssertionError('even_unselected_nonfinite_branches_must_refuse')
    model.load_state_dict(state)
    with torch.no_grad():
        model.log_gate.bias.fill_(100)
    assert model(torch.tensor([[2., -1., 2.]]), torch.ones(1, 3), [(1, 3)]).tolist() == [-2.]
    parent = WindowGateModel(Encoder())
    extra = {'log_head.weight', 'log_gate.weight', 'log_gate.bias'}
    assert set(state) - set(parent.state_dict()) == extra
    encoder = Encoder(); encoder.config = SimpleNamespace(hidden_size=768)
    assert sum(v.numel() for k, v in SourceKindGateModel(encoder).state_dict().items() if k in extra) == 1537


if __name__ == '__main__':
    self_check()
    print('Source-kind boundaries, service priority, pooling, extra state and nonfinite refusal pass.')