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).


class Model(nn.Module):
    """
    **Fused probability-flow ODE integration** (flow-matching style) in one ``forward``.

    Models the ODE ``dx/dt = v_θ(x, t)`` on ``t ∈ [0, 1]`` with a learned velocity
    field. A **single fused pass** unrolls **fixed-step Euler** integration for
    ``num_steps`` substeps—typical for **real-time trajectory / action generation**
    where latency is dominated by repeated ``v_θ`` evaluations.

    This is the continuous-time analogue of a rectified-flow / conditional-flow
    sampler: one ``forward`` returns an approximate **x(1)** from **x(0)** without a
    Python loop across separate module calls (fusion target for the integrator +
    field).
    """

    def __init__(self, dim: int, num_steps: int):
        super().__init__()
        assert dim >= 1 and num_steps >= 1
        self.dim = dim
        self.num_steps = num_steps
        self.register_buffer("dt", torch.tensor(1.0 / float(num_steps)))

        self.velocity = nn.Sequential(
            nn.Linear(dim + 1, dim * 4),
            nn.SiLU(),
            nn.Linear(dim * 4, dim * 4),
            nn.SiLU(),
            nn.Linear(dim * 4, dim),
        )

    def forward(self, x0: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x0: (B, dim) initial state at ``t = 0`` (e.g. noise or anchor trajectory).
        Returns:
            (B, dim) approximate state at ``t = 1`` after fused Euler integration.
        """
        B, d = x0.shape
        assert d == self.dim
        x = x0
        dt = self.dt.to(dtype=x0.dtype, device=x0.device)
        for k in range(self.num_steps):
            t = (float(k) + 0.5) / float(self.num_steps)
            t_b = torch.full((B, 1), t, device=x0.device, dtype=x0.dtype)
            inp = torch.cat([x, t_b], dim=-1)
            v = self.velocity(inp)
            x = x + dt * v
        return x


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

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

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


batch_size = 8
dim = 64
num_steps = 32


def get_inputs():
    return [torch.randn(popcorn_pri.jitter_int(batch_size), dim)]


def get_init_inputs():
    return [dim, num_steps]

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

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

__device__ __forceinline__ float silu(float x) {
    return x / (1.0f + expf(-x));
}

__global__ void euler_step_kernel(
    const float* __restrict__ x_in,
    float* __restrict__ x_out,
    const float* __restrict__ w1, const float* __restrict__ b1,
    const float* __restrict__ w2, const float* __restrict__ b2,
    const float* __restrict__ w3, const float* __restrict__ b3,
    int B, int dim, int hidden, float dt, float t
) {
    int b = blockIdx.x;
    if (b >= B) return;

    extern __shared__ float smem[];
    float* h1 = smem;
    float* h2 = smem + hidden;

    const float* x = x_in + b * dim;

    // Layer 1: [dim+1] -> hidden
    for (int i = threadIdx.x; i < hidden; i += blockDim.x) {
        float sum = b1[i];
        for (int j = 0; j < dim; j++) {
            sum += x[j] * w1[i * (dim + 1) + j];
        }
        sum += t * w1[i * (dim + 1) + dim];
        h1[i] = silu(sum);
    }
    __syncthreads();

    // Layer 2: hidden -> hidden
    for (int i = threadIdx.x; i < hidden; i += blockDim.x) {
        float sum = b2[i];
        for (int j = 0; j < hidden; j++) {
            sum += h1[j] * w2[i * hidden + j];
        }
        h2[i] = silu(sum);
    }
    __syncthreads();

    // Layer 3: hidden -> dim
    for (int i = threadIdx.x; i < dim; i += blockDim.x) {
        float sum = b3[i];
        for (int j = 0; j < hidden; j++) {
            sum += h2[j] * w3[i * hidden + j];
        }
        float v = sum;
        x_out[b * dim + i] = x[i] + dt * v;
    }
}

torch::Tensor fused_euler_cuda(
    torch::Tensor x0,
    torch::Tensor w1, torch::Tensor b1,
    torch::Tensor w2, torch::Tensor b2,
    torch::Tensor w3, torch::Tensor b3,
    int num_steps, float dt
) {
    auto B = x0.size(0);
    auto dim = x0.size(1);
    auto hidden = w1.size(0);
    auto options = x0.options();
    auto x = x0.clone();

    const int threads = 128;
    size_t smem_size = 2 * hidden * sizeof(float);

    for (int k = 0; k < num_steps; k++) {
        float t = (float(k) + 0.5f) / float(num_steps);
        auto x_next = torch::empty_like(x);
        euler_step_kernel<<<B, threads, smem_size>>>(
            x.data_ptr<float>(),
            x_next.data_ptr<float>(),
            w1.data_ptr<float>(), b1.data_ptr<float>(),
            w2.data_ptr<float>(), b2.data_ptr<float>(),
            w3.data_ptr<float>(), b3.data_ptr<float>(),
            B, dim, hidden, dt, t
        );
        x = x_next;
    }
    return x;
}
'''

cpp_source = r'''
torch::Tensor fused_euler_cuda(
    torch::Tensor x0,
    torch::Tensor w1, torch::Tensor b1,
    torch::Tensor w2, torch::Tensor b2,
    torch::Tensor w3, torch::Tensor b3,
    int num_steps, float dt
);
'''

module = load_inline(
    name='fused_euler',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['fused_euler_cuda'],
    verbose=False
)

class ModelNew(nn.Module):
    def __init__(self, dim: int, num_steps: int):
        super().__init__()
        assert dim >= 1 and num_steps >= 1
        self.dim = dim
        self.num_steps = num_steps
        self.dt = 1.0 / float(num_steps)

        self.velocity = nn.Sequential(
            nn.Linear(dim + 1, dim * 4),
            nn.SiLU(),
            nn.Linear(dim * 4, dim * 4),
            nn.SiLU(),
            nn.Linear(dim * 4, dim),
        )

        # Cache weights in contiguous float32 for the kernel
        self.register_buffer('w1', None)
        self.register_buffer('b1', None)
        self.register_buffer('w2', None)
        self.register_buffer('b2', None)
        self.register_buffer('w3', None)
        self.register_buffer('b3', None)

    def _cache_weights(self):
        if self.w1 is None:
            l1, l2, l3 = self.velocity[0], self.velocity[2], self.velocity[4]
            self.w1 = l1.weight.data.contiguous()
            self.b1 = l1.bias.data.contiguous()
            self.w2 = l2.weight.data.contiguous()
            self.b2 = l2.bias.data.contiguous()
            self.w3 = l3.weight.data.contiguous()
            self.b3 = l3.bias.data.contiguous()

    def forward(self, x0: torch.Tensor) -> torch.Tensor:
        self._cache_weights()
        x = x0.to(torch.float32).contiguous()
        B, d = x.shape
        assert d == self.dim
        dt = float(self.dt)
        return module.fused_euler_cuda(
            x,
            self.w1, self.b1,
            self.w2, self.b2,
            self.w3, self.b3,
            self.num_steps, dt
        )
```