"""Prototype: frozen-parent encoder/base head with a learned CLS-window selector."""
import torch
from torch import nn
from tokenfold_select._model import JointGroupModel


class WindowGateModel(JointGroupModel):
    """One head choice per complete encoder window, not per token or group."""
    def __init__(self, encoder):
        super().__init__(encoder)
        self.window_gate=nn.Linear(encoder.config.hidden_size,1)
        self.service_head=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)
        valid=torch.isfinite(gate)&torch.isfinite(original)&torch.isfinite(service)
        chosen=torch.where(gate>=0,service,original)
        return torch.where(valid,chosen,torch.full_like(chosen,float('nan')))

    def forward(self, input_ids, attention_mask, ranges):
        if input_ids.ndim!=2 or input_ids.shape[0]!=1 or attention_mask.shape!=input_ids.shape:
            raise ValueError('single_sequence_required')
        if not 1<=len(ranges)<=512:
            raise ValueError('invalid_ranges')
        previous=0
        for start,end in ranges:
            if not previous<=start<end<=input_ids.shape[1] or not attention_mask[0,start:end].all():
                raise ValueError('invalid_ranges')
            previous=end
        positions=torch.arange(input_ids.shape[1],device=input_ids.device).unsqueeze(0)
        states=self.encoder(input_ids=input_ids,attention_mask=attention_mask,position_ids=positions).last_hidden_state[0]
        pooled=torch.stack([states[a:b].mean(0) for a,b in ranges])
        scores=self.choose_head(pooled,states[:1]).squeeze(-1)
        if not torch.isfinite(scores).all():
            raise ValueError('nonfinite_scores')
        return scores


def self_check():
    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))
    m=WindowGateModel(Encoder()).eval()
    with torch.no_grad():
        m.head.weight.fill_(10);m.head.bias.zero_()
        m.service_head.weight.fill_(-10);m.service_head.bias.zero_()
        m.window_gate.weight.fill_(1);m.window_gate.bias.zero_()
    for first,expected in [(-2.,5.),(2.,-5.),(0.,-5.)]:
        x=torch.tensor([[first,-1.,2.]])
        result=m(x,torch.ones_like(x),[(1,3)])
        assert result.tolist()==[expected]
        states=x.unsqueeze(-1)
        tokens=m.choose_head(states,states[:,:1,:]).squeeze(-1)
        assert torch.equal(tokens[:,1:3].mean(1),result)
    with torch.no_grad():m.service_head.weight.fill_(float('inf'))
    try:m(torch.tensor([[-2.,-1.,2.]]),torch.ones(1,3),[(1,3)])
    except ValueError as error:assert str(error)=='nonfinite_scores'
    else:raise AssertionError('masked nonfinite branch must refuse, not fall back')


if __name__=='__main__':self_check();print('CLS-global head choice, pooling parity, zero-boundary and nonfinite refusal checks passed.')