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_runtime.h>
#include <cmath>

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

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

    for (int k = 0; k < num_steps; k++) {
        float t = (float(k) + 0.5f) / float(num_steps);
        float h1[256];
        for (int h = 0; h < hidden; h++) {
            float sum = b1[h];
            for (int i = 0; i < dim; i++) {
                sum += W1[h * (dim + 1) + i] * x[i];
            }
            sum += W1[h * (dim + 1) + dim] * t;
            h1[h] = sum / (1.0f + expf(-sum));
        }
        float h2[256];
        for (int h = 0; h < hidden; h++) {
            float sum = b2[h];
            for (int i = 0; i < hidden; i++) {
                sum += W2[h * hidden + i] * h1[i];
            }
            h2[h] = sum / (1.0f + expf(-sum));
        }
        float v[64];
        for (int d = 0; d < dim; d++) {
            float sum = b3[d];
            for (int i = 0; i < hidden; i++) {
                sum += W3[d * hidden + i] * h2[i];
            }
            v[d] = sum;
        }
        for (int i = 0; i < dim; i++) {
            x[i] += dt * v[i];
        }
    }

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

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

    dim3 block(32);
    dim3 grid((B + 31) / 32);

    euler_integrate_kernel<<<grid, block>>>(
        x0.data_ptr<float>(),
        x_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, dt, num_steps
    );

    return x_out;
}
"""

cpp_source = """
torch::Tensor euler_integrate_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
);
"""

class ModelNew(nn.Module):
    def __init__(self, dim: int, num_steps: int):
        super().__init__()
        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),
        )
        self._cuda_module = load_inline(
            name="euler_integrate",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["euler_integrate_cuda"],
            verbose=False,
            extra_cuda_cflags=['-O3']
        )

    def forward(self, x0: torch.Tensor) -> torch.Tensor:
        W1 = self.velocity[0].weight
        b1 = self.velocity[0].bias
        W2 = self.velocity[2].weight
        b2 = self.velocity[2].bias
        W3 = self.velocity[4].weight
        b3 = self.velocity[4].bias
        return self._cuda_module.euler_integrate_cuda(
            x0, W1, b1, W2, b2, W3, b3, self.num_steps, self.dt
        )