"""Explicit native-math primitives; CPU remains a mathematical diagnostic."""
from functools import lru_cache
from pathlib import Path
import torch
from .kernel_build import stable_source


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


def unary(value,operation):
    value=value.float().contiguous()
    if value.is_cuda:
        if torch.cuda.get_device_capability(value.device)<(7,5):
            raise RuntimeError('Native tanh.approx reference currently requires SM75+; older-SM emulation remains unvalidated')
        return extension().unary(value,operation)
    functions=(torch.tanh,torch.tanh,torch.cos,torch.sin,torch.exp,torch.sigmoid)
    return functions[operation](value)


def bf16_product(a,b,accumulator=None):
    """BF16 operands with FP32 accumulation; no intermediate BF16 output cast."""
    if a.dtype!=torch.bfloat16 or b.dtype!=torch.bfloat16:
        raise ValueError('BF16 matrix operands required')
    shape=(*a.shape[:-2],a.shape[-2],b.shape[-1])
    if a.is_cuda:
        left=a.reshape(-1,*a.shape[-2:]).contiguous()
        right=b.reshape(-1,*b.shape[-2:]).contiguous()
        if torch.cuda.get_device_capability(a.device)<(8,0):
            raise RuntimeError('Native BF16 MMA reference requires SM80+; older-SM emulation remains unvalidated')
        initial=None if accumulator is None else accumulator.reshape(-1,*shape[-2:]).contiguous()
        canonical=torch.cuda.get_device_capability(a.device)!=(10,0)
        return extension().mma(left,right,initial,canonical).reshape(shape)
    value=a.float()@b.float()
    return value if accumulator is None else accumulator+value


def cuda_fma(a,b,c):
    a,b,c=torch.broadcast_tensors(a.float(),b.float(),c.float())
    return extension().fma(a.contiguous(),b.contiguous(),c.contiguous())


def f32_batched_product(a,b,*,streams,stride):
    if a.dtype!=torch.float32 or b.dtype!=torch.float32 or a.shape[:-2]!=b.shape[:-2]:
        raise ValueError('Matching FP32 batch dimensions required')
    shape=(*a.shape[:-2],a.shape[-2],b.shape[-1])
    left=a.reshape(-1,*a.shape[-2:]).contiguous();right=b.reshape(-1,*b.shape[-2:]).contiguous()
    return extension().bmm_reference(left,right,streams,stride).reshape(shape)