"""Bounded weight preparation and executable current-recipe linears."""
from functools import lru_cache
from pathlib import Path
import torch
from .kernel_build import stable_source
import torch.nn.functional as F
from .numerics import Packed,quantize,bridge_quantize,gemm_release_reference
from .contracts import blend

@lru_cache(None)
def extension():
    from torch.utils.cpp_extension import load
    return load(name='fidel_compatibility_packed_reference_v1',
                sources=[str(stable_source(Path(__file__).with_name('packed_gemm_reference.cu')))],
                extra_cuda_cflags=['-O3','--fmad=false'],verbose=False)

def packed_product(a,b):
    if a.data.is_cuda:
        return extension().gemm(a.data,b.data,a.scales,b.scales,a.global_scale,b.global_scale)
    return gemm_release_reference(a,b)

def bridge_values(x):
    if x.is_cuda:return extension().bridge(x.float().contiguous())
    return bridge_quantize(x)

# Native B200 FP32 dense profiles under the release's fixed 128-row baseline.
# Keys are (output width, reduction width); every admitted profile has retained
# real-operand authority. Other shapes remain behind the model fidelity gate.
DENSE_F32_PARTITIONS={
    (16,1536):24,(32,1536):24,(384,1536):24,(1536,1536):24,
    (512,1536):16,(256,1536):16,(12288,1536):1,(1536,6144):24,
    (1536,512):1,(1536,256):1,(1536,32):1,(384,32):1,
    (128,1536):24,(64,1536):24,(1536,128):1,(1536,64):1,
    (131072,1536):1,(131072,256):1,
}

class LinearEngine:
    def __init__(self,weight_rows=512):
        self.weight_rows=weight_rows
        self.weights={}
        self.dense_weights={}
        self.calls=0
    def clear(self):
        self.weights.clear();self.dense_weights.clear()
    def prepare(self,name,weight,seed,bridge):
        key=(name,seed,bridge,weight.device,weight.dtype,weight.data_ptr(),weight._version,tuple(weight.stride()))
        if key in self.weights:return self.weights[key]
        global_scale=(weight.float().abs().amax()*(1./(6*448*(16/17)))).clamp_min(torch.finfo(torch.float32).tiny)
        data,scales,original,means=[],[],[],[]
        for start in range(0,weight.shape[0],self.weight_rows):
            block=weight[start:start+self.weight_rows].float().contiguous()
            q=quantize(block,eden=True,seed=seed+1,global_scale_override=global_scale,
                       group_offset=start*(weight.shape[1]//16))
            data.append(q.data);scales.append(q.scales)
            if bridge:
                qb,sb=bridge_values(block)
                original.append(qb);means.append(sb.mean(1,keepdim=True))
        value=(Packed(torch.cat(data),torch.cat(scales),global_scale),
               torch.cat(original) if bridge else None,torch.cat(means) if bridge else None)
        self.weights[key]=value
        return value
    def __call__(self,x,weight,bias=None,*,name,seed=0,bridge=False,activation_row_offset=0):
        if x.requires_grad or weight.requires_grad:raise ValueError('Inference only')
        if x.shape[-1]!=weight.shape[1]:raise ValueError('Linear width mismatch')
        if type(activation_row_offset) is not int or activation_row_offset<0 or activation_row_offset%128:
            raise ValueError('Activation offsets must start on a 128-row tile')
        self.calls+=1
        shape=x.shape;rows=x.reshape(-1,shape[-1]);live=rows.shape[0]
        padded=F.pad(rows,(0,0,0,(-live)%128))
        qw,bridge_w,bridge_mean=self.prepare(name,weight,seed,bridge)
        qx=quantize(padded.float(),eden=True,seed=seed,rowwise=True,
                    group_offset=activation_row_offset*(padded.shape[1]//16))
        recipe=packed_product(qx,qw)
        if bias is not None:recipe=recipe+bias.float()
        recipe=recipe[:live].to(x.dtype)
        output=[]
        for start in range(0,live,128):
            block=padded[start:start+128];count=min(128,live-start)
            if bridge:
                q,s=bridge_values(block)
                dot=(q@bridge_w.T).to(torch.bfloat16)
                mean=s.mean(1,keepdim=True)
                baseline=(dot*mean)*bridge_mean.T if x.dtype==torch.bfloat16 else dot.float()*mean.float()*bridge_mean.T.float()
                if bias is not None:baseline=baseline+bias.to(baseline.dtype)
            elif block.is_cuda and block.dtype==weight.dtype==torch.bfloat16 and torch.cuda.get_device_capability(block.device)!=(10,0):
                from .math_primitives import bf16_product
                key=(weight.data_ptr(),weight._version,tuple(weight.shape),tuple(weight.stride()),str(weight.device))
                if key not in self.dense_weights:self.dense_weights[key]=weight.T.contiguous()
                baseline=bf16_product(block[:((count+15)//16)*16],self.dense_weights[key]).to(block.dtype)
                if bias is not None:baseline=baseline+bias.to(baseline.dtype)
            elif block.is_cuda and block.dtype==torch.float32 and torch.cuda.get_device_capability(block.device)!=(10,0) and tuple(weight.shape) in DENSE_F32_PARTITIONS:
                from .math_primitives import extension as math_extension
                partitions=DENSE_F32_PARTITIONS[tuple(weight.shape)]
                baseline=math_extension().dense_fma(block[:count].float().contiguous(),weight.float().contiguous(),partitions)
                if bias is not None:baseline=baseline+bias.float()
            elif block.dtype==weight.dtype:
                baseline=F.linear(block,weight,bias)
            else:
                baseline=F.linear(block.float(),weight.float(),None if bias is None else bias.float()).to(x.dtype)
            output.append(blend(baseline[:count],recipe[start:start+count],legacy_bridge=bridge))
        return torch.cat(output).reshape(*shape[:-1],weight.shape[0])