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 = '''
#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 fused_ode_kernel(
    const float* __restrict__ x0,
    float* __restrict__ 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, int num_steps, float dt
) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;

    float x[64];  // assume dim <=64
    for (int i = 0; i < dim; i++) {
        x[i] = x0[b * dim + i];
    }

    float t_step = 1.0f / num_steps;
    for (int k = 0; k < num_steps; k++) {
        float t = (k + 0.5f) * t_step;

        // inp: x + t , size dim+1
        // layer1: hidden = dim*4
        float h1[256];
        for (int h = 0; h < hidden; h++) {
            float sum = b1[h];
            for (int i = 0; i < dim; i++) {
                sum += x[i] * w1[h * (dim+1) + i];
            }
            sum += t * w1[h * (dim+1) + dim];
            h1[h] = silu(sum);
        }

        // layer2
        float h2[256];
        for (int h = 0; h < hidden; h++) {
            float sum = b2[h];
            for (int i = 0; i < hidden; i++) {
                sum += h1[i] * w2[h * hidden + i];
            }
            h2[h] = silu(sum);
        }

        // layer3
        float v[64];
        for (int d = 0; d < dim; d++) {
            float sum = b3[d];
            for (int i = 0; i < hidden; i++) {
                sum += h2[i] * w3[d * hidden + i];
            }
            v[d] = sum;
        }

        for (int d = 0; d < dim; d++) {
            x[d] += dt * v[d];
        }
    }

    for (int i = 0; i < dim; i++) {
        out[b * dim + i] = x[i];
    }
}

torch::Tensor fused_ode(
    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_val
) {
    int B = x0.size(0);
    int dim = x0.size(1);
    int hidden = dim * 4;

    auto out = torch::empty_like(x0);

    dim3 blocks((B + 255) / 256);
    dim3 threads(256);

    fused_ode_kernel<<<blocks, threads>>>(
        x0.data_ptr<float>(),
        out.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, num_steps, dt_val
    );

    return out;
}
'''

cpp_source = '''
torch::Tensor fused_ode(
    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_val
);
'''

fused_module = load_inline(
    name='fused_ode_module',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['fused_ode'],
    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),
        )

    def forward(self, x0: torch.Tensor) -> torch.Tensor:
        if not x0.is_cuda or x0.dtype != torch.float32:
            x0 = x0.cuda().float()
        
        # Get weights
        lin1 = self.velocity[0]
        lin2 = self.velocity[2]
        lin3 = self.velocity[4]
        
        w1, b1 = lin1.weight, lin1.bias
        w2, b2 = lin2.weight, lin2.bias
        w3, b3 = lin3.weight, lin3.bias
        
        return fused_module.fused_ode(
            x0, w1, b1, w2, b2, w3, b3,
            self.num_steps, self.dt
        )
```