"""Owner-selected bf16 segmented attention and shared-prefix inference."""
import os
import torch
import saif_v2_core as core
from transformers.models.qwen3_5 import modeling_qwen3_5 as modeling

MASK=core.attention_mask
ATTENTION=modeling.ALL_ATTENTION_FUNCTIONS['sdpa']
LAYOUT=None


def mask(layout, attention, device, dtype, boolean=False):
    global LAYOUT
    assert attention=='bidir'
    LAYOUT=layout
    return {'full_attention':None,'linear_attention':None}


def segmented(module,query,key,value,attention_mask,dropout=0.,scaling=None,**kwargs):
    assert dropout==0 and LAYOUT is not None
    prefix=LAYOUT.prefix
    assert query.shape[-2]==len(LAYOUT.ids)
    if module.num_key_value_groups>1:
        key=key.repeat_interleave(module.num_key_value_groups,dim=1)
        value=value.repeat_interleave(module.num_key_value_groups,dim=1)
    def attend(q,k,v):
        return torch.nn.functional.scaled_dot_product_attention(q,k,v,dropout_p=0.,is_causal=False,scale=scaling)
    output=[attend(query[:,:,:prefix],key[:,:,:prefix],value[:,:,:prefix])]
    for start,end in LAYOUT.blocks:
        output.append(attend(query[:,:,start:end],torch.cat((key[:,:,:prefix],key[:,:,start:end]),dim=2),torch.cat((value[:,:,:prefix],value[:,:,start:end]),dim=2)))
    return torch.cat(output,dim=2).transpose(1,2).contiguous(),None


def install():
    core.attention_mask=mask
    modeling.ALL_ATTENTION_FUNCTIONS['sdpa']=segmented
    os.environ['SAIF_EXPERIMENTAL_SHARED_DELTA']='1'


def uninstall():
    core.attention_mask=MASK
    modeling.ALL_ATTENTION_FUNCTIONS['sdpa']=ATTENTION
    os.environ.pop('SAIF_EXPERIMENTAL_SHARED_DELTA',None)