import torch
import torch.nn as nn

# Default ModelNew is PyTorch-only. run_and_check: use check_kernel=False with backend=cuda
# until the candidate file contains a real CUDA kernel (see kernel_static_checker).


def _fake_quant_symmetric_int8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    """Symmetric int8 fake-quant: round(x/scale) clamped to [-127, 127], times scale."""
    q = torch.clamp(torch.round(x / scale), -127.0, 127.0)
    return q * scale


class Model(nn.Module):
    """
    **QuantVLA selective layout fusion** (reference target for a fused CUDA/Triton kernel).

    Vision–language–action stacks often use **token-wise mixed precision**: some
    timesteps or modalities tolerate aggressive quantization while others stay in
    full precision. This module fuses in **one forward**:

      1. **Layout routing** — ``sigmoid(layout_logits)`` yields a per-token mix
         coefficient in ``(0, 1)`` (quant vs FP residual).
      2. **Quantitative path** — global symmetric **int8-style fake quantization**
         (scale from the tensor max, ``/ 127``).
      3. **Selective blend** — ``α · x_quant + (1 - α) · x``.
      4. **Head** — linear projection on the mixed activations.

    This matches a “quantitative selective layout” block: the layout logits (from a
    router, modality tags, or calibration) choose how strongly each position uses the
    quantized layout vs the identity layout.
    """

    def __init__(self, dim: int):
        super().__init__()
        self.dim = dim
        self.out_proj = nn.Linear(dim, dim, bias=True)

    def forward(self, x: torch.Tensor, layout_logits: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, T, D) VLA token activations (fused modalities).
            layout_logits: (B, T) unbounded logits; larger → more mass on quantized path.
        Returns:
            (B, T, D) updated activations.
        """
        B, T, D = x.shape
        assert D == self.dim
        assert layout_logits.shape == (B, T)

        alpha = torch.sigmoid(layout_logits).unsqueeze(-1)
        with torch.no_grad():
            amax = x.detach().abs().max().clamp(min=1e-8)
        scale = amax / 127.0
        x_q = _fake_quant_symmetric_int8(x, scale)
        mixed = alpha * x_q + (1.0 - alpha) * x
        return self.out_proj(mixed)


class ModelNew(nn.Module):
    """KernelBench candidate entry point; default delegates to ``Model``."""

    def __init__(self, dim: int):
        super().__init__()
        self._impl = Model(dim)

    def forward(self, x: torch.Tensor, layout_logits: torch.Tensor) -> torch.Tensor:
        return self._impl(x, layout_logits)


batch_size = 8
seq_len = 64
dim = 128


def get_inputs():
    p = popcorn_pri
    x = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), dim)
    layout_logits = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len))
    return [x, layout_logits]


def get_init_inputs():
    return [dim]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cuda_source = '''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

__global__ void compute_mixed_kernel(const float* __restrict__ x,
                                     const float* __restrict__ logits,
                                     float* __restrict__ mixed,
                                     int B, int T, int D, float scale) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total_tokens = B * T;
    if (idx >= total_tokens) return;

    float alpha = 1.0f / (1.0f + __expf(-logits[idx]));
    float one_minus_alpha = 1.0f - alpha;
    const float* x_row = x + idx * D;
    float* mixed_row = mixed + idx * D;

    for (int d = 0; d < D; d++) {
        float val = x_row[d];
        float q = roundf(val / scale);
        q = fmaxf(fminf(q, 127.0f), -127.0f);
        float xq = q * scale;
        mixed_row[d] = alpha * xq + one_minus_alpha * val;
    }
}

torch::Tensor compute_mixed_cuda(torch::Tensor x, torch::Tensor logits, float scale) {
    auto B = x.size(0);
    auto T = x.size(1);
    auto D = x.size(2);
    auto mixed = torch::empty_like(x);

    int total_tokens = B * T;
    int threads = 256;
    int blocks = (total_tokens + threads - 1) / threads;

    compute_mixed_kernel<<<blocks, threads>>>(
        x.data_ptr<float>(),
        logits.data_ptr<float>(),
        mixed.data_ptr<float>(),
        B, T, D, scale
    );

    return mixed;
}
'''

cpp_source = '''
torch::Tensor compute_mixed_cuda(torch::Tensor x, torch::Tensor logits, float scale);
'''

# Compile the extension
custom_ops = load_inline(
    name='custom_mixed_op',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['compute_mixed_cuda'],
    verbose=False,
    extra_cuda_cflags=['-O3']
)

class ModelNew(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.dim = dim
        self.out_proj = nn.Linear(dim, dim, bias=True)

    def forward(self, x: torch.Tensor, layout_logits: torch.Tensor) -> torch.Tensor:
        B, T, D = x.shape
        assert D == self.dim
        assert layout_logits.shape == (B, T)

        with torch.no_grad():
            amax = x.detach().abs().max().clamp(min=1e-8)
        scale = amax / 127.0

        mixed = custom_ops.compute_mixed_cuda(x, layout_logits, scale.item())
        return self.out_proj(mixed)
```