"""Fuse equal-rotation packed QKV and gate/up projections without BF16 weights."""
import types
import torch
from torch import nn
from orbitquant import OrbitQuantConfig
from orbitquant.layers import OrbitQuantLinear


class PackedSlice(nn.Module):
    """Compatibility view; the owning attention/MLP registers the packed group."""
    def __init__(self, group, start, length):
        super().__init__()
        object.__setattr__(self, '_group', group)
        self.start, self.out_features, self.in_features = start, length, group.in_features
        self.bias = None
    def forward(self, x):
        return self._group(x)[..., self.start:self.start+self.out_features]


def group(modules, config, name):
    first=modules[0]
    if not all(isinstance(m,OrbitQuantLinear) and m.in_features==first.in_features and m.bias is None for m in modules):
        raise ValueError('Fusion requires compatible bias-free packed linears')
    for m in modules[1:]:
        if (m.weight_bits,m.activation_bits,m.activation_eps,m.rotation.seed,m.rotation.block_size)!=(first.weight_bits,first.activation_bits,first.activation_eps,first.rotation.seed,first.rotation.block_size):
            raise ValueError('Fusion requires identical quantization and rotation')
    return OrbitQuantLinear(in_features=first.in_features,out_features=sum(m.out_features for m in modules),
        config=config,module_name=name,source_weight_layout='out_in',bias=None,
        packed_weight_indices=torch.cat([m.packed_weight_indices for m in modules]),
        row_norms=torch.cat([m.row_norms for m in modules]),debug_weight=None).eval()


def project_qkv(self,x,cos,sin):
    from yue2.modeling_yue2 import _apply_rotary
    b,t,_=x.shape
    q,k,v=self.packed_qkv(x).split((self.num_heads*self.head_dim,self.num_kv_heads*self.head_dim,self.num_kv_heads*self.head_dim),dim=-1)
    q=self.q_norm(q.reshape(b,t,self.num_heads,self.head_dim))
    k=self.k_norm(k.reshape(b,t,self.num_kv_heads,self.head_dim))
    return _apply_rotary(q,cos.unsqueeze(2),sin.unsqueeze(2)), _apply_rotary(k,cos.unsqueeze(2),sin.unsqueeze(2)), v.reshape(b,t,self.num_kv_heads,self.head_dim)


def mlp_forward(self,x):
    gate,up=self.packed_gate_up(x).chunk(2,dim=-1)
    return self.down_proj(torch.nn.functional.silu(gate)*up)


def fuse_model(model,manifest):
    if getattr(model,'_orbitquant_fused',False): return model
    config=OrbitQuantConfig.from_dict(manifest['config'])
    for i,layer in enumerate(model.model.layers):
        for attr in ('self_attn','nar_self_attn'):
            attn=getattr(layer,attr)
            parts=[getattr(attn,n) for n in ('q_proj','k_proj','v_proj')]
            packed=group(parts,config,f'model.layers.{i}.{attr}.packed_qkv')
            attn.packed_qkv=packed
            offset=0
            for name,part in zip(('q_proj','k_proj','v_proj'),parts):
                setattr(attn,name,PackedSlice(packed,offset,part.out_features));offset+=part.out_features
            attn.project_qkv=types.MethodType(project_qkv,attn)
            del parts
        for attr in ('mlp','nar_mlp'):
            mlp=getattr(layer,attr)
            n=mlp.gate_proj.out_features
            packed=group([mlp.gate_proj,mlp.up_proj],config,f'model.layers.{i}.{attr}.packed_gate_up')
            mlp.packed_gate_up=packed
            mlp.gate_proj=PackedSlice(packed,0,n);mlp.up_proj=PackedSlice(packed,n,n)
            mlp.forward=types.MethodType(mlp_forward,mlp)
    model._orbitquant_fused=True
    return model