"""Fixed-length, constrained HF beam search and independent numerical checks."""
import torch
from transformers import GenerationConfig, LogitsProcessor, LogitsProcessorList

class VisualTokensOnly(LogitsProcessor):
    def __init__(self, valid_size):
        self.valid_size = valid_size

    def __call__(self, input_ids, scores):
        scores = scores.clone()
        scores[:, self.valid_size:] = -torch.inf
        return scores

@torch.inference_mode()
def beam_decode(model, prefix, width=4, steps=80, valid_size=4375):
    config = GenerationConfig(
        do_sample=False, num_beams=width, num_return_sequences=width,
        max_new_tokens=steps, min_new_tokens=steps,
        eos_token_id=None, bos_token_id=None, pad_token_id=0,
        length_penalty=0.0 if width > 1 else 1.0, early_stopping=False, renormalize_logits=True,
        repetition_penalty=1.0, no_repeat_ngram_size=0,
        use_cache=True, return_dict_in_generate=True, output_scores=True,
    )
    result = model.generate(
        input_ids=prefix, attention_mask=torch.ones_like(prefix),
        generation_config=config,
        logits_processor=LogitsProcessorList([VisualTokensOnly(valid_size)]),
        eos_token_id=None, bos_token_id=None,
        forced_bos_token_id=None, forced_eos_token_id=None,
    )
    candidates = result.sequences[:, prefix.shape[1]:]
    assert candidates.shape == (width,steps)
    assert int(candidates.min()) >= 0 and int(candidates.max()) < valid_size
    assert len({tuple(x) for x in candidates.tolist()}) == width
    transitions = model.compute_transition_scores(
        result.sequences, result.scores, getattr(result,'beam_indices',None),
        normalize_logits=False,
    )
    assert transitions.shape == (width,steps) and torch.isfinite(transitions).all()
    sums = transitions.sum(-1)
    scores = result.sequences_scores if width > 1 else sums
    error = float((sums-scores).abs().max())
    assert torch.allclose(sums,scores,rtol=0,atol=2e-4), error
    assert torch.all(scores[:-1] >= scores[1:])
    return candidates, scores, {'beam_width':width,'tokens_per_frame':steps,
        'selected_index':0,'candidate_scores':scores.tolist(),
        'score_reconstruction_max_error':error,
        'score_definition':'Sum of log probabilities renormalized over valid visual tokens only; no length penalty or GT/image reward.'}

@torch.inference_mode()
def independent_check():
    from transformers import LlamaConfig, LlamaForCausalLM
    torch.manual_seed(3719)
    model = LlamaForCausalLM(LlamaConfig(vocab_size=8,hidden_size=32,
        intermediate_size=48,num_hidden_layers=2,num_attention_heads=4,
        num_key_value_heads=4,max_position_embeddings=128,
        eos_token_id=2,bos_token_id=1,pad_token_id=0,
        attention_dropout=0.0)).float().eval()
    prefix = torch.tensor([[6,1,3]],dtype=torch.long)
    reports=[]
    for width in [1,4]:
        beams=[([],0.0)]
        for _ in range(4):
            expanded=[]
            for seq,score in beams:
                x=torch.tensor([prefix[0].tolist()+seq],dtype=torch.long)
                probs=model(input_ids=x,use_cache=False).logits[0,-1,:5].log_softmax(-1)
                expanded.extend((seq+[token],score+float(probs[token])) for token in range(5))
            beams=sorted(expanded,key=lambda item:item[1],reverse=True)[:width]
        ids,scores,meta=beam_decode(model,prefix,width=width,steps=4,valid_size=5)
        expected=torch.tensor([seq for seq,_ in beams],dtype=torch.long)
        expected_scores=torch.tensor([score for _,score in beams])
        assert torch.equal(ids,expected),(width,ids,expected)
        assert torch.allclose(scores,expected_scores,rtol=0,atol=2e-5)
        reports.append({'width':width,'matches_independent_uncached_search':True,
            'max_score_error':float((scores-expected_scores).abs().max())})
    return {'tests':reports,'eos_disabled':True,'valid_token_constraint_checked':True,'complete':True}

if __name__ == '__main__':
    import json, transformers
    print(json.dumps({'transformers':transformers.__version__,**independent_check()}))