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

_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor flow_forward(torch::Tensor x0,
                           torch::Tensor w1, torch::Tensor b1,
                           torch::Tensor w2, torch::Tensor b2,
                           torch::Tensor w3, torch::Tensor b3,
                           int64_t num_steps);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be CUDA")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")

__device__ __forceinline__ float silu_f(float x) {
    return x / (1.0f + __expf(-x));
}

__global__ __launch_bounds__(1024, 1)
void flow_kernel_64_4way(const float* __restrict__ x0,
                         const float* __restrict__ w1,
                         const float* __restrict__ b1,
                         const float* __restrict__ w2,
                         const float* __restrict__ b2,
                         const float* __restrict__ w3,
                         const float* __restrict__ b3,
                         float* __restrict__ out,
                         int B, int num_steps) {
    const int D = 64;
    const int H = 256;
    int b = blockIdx.x;
    int tid = threadIdx.x;

    __shared__ float x[D];
    __shared__ float h1[H];
    __shared__ float h2[H];
    __shared__ float part[1024];

    if (tid < D) x[tid] = x0[b * D + tid];
    __syncthreads();

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

        // First linear: four threads cooperatively compute each of 256 neurons.
        int j = tid >> 2;
        int lane = tid & 3;
        float s = 0.0f;
        #pragma unroll
        for (int i = lane; i < D; i += 4) {
            s = fmaf(x[i], w1[j * (D + 1) + i], s);
        }
        if (lane == 0) s = fmaf(t, w1[j * (D + 1) + D], s);
        part[tid] = s;
        __syncthreads();
        if (lane == 0) {
            int base = j << 2;
            float v = b1[j] + part[base] + part[base + 1] + part[base + 2] + part[base + 3];
            h1[j] = silu_f(v);
        }
        __syncthreads();

        // Second linear: four-way split for each hidden neuron.
        s = 0.0f;
        #pragma unroll 4
        for (int i = lane; i < H; i += 4) {
            s = fmaf(h1[i], w2[j * H + i], s);
        }
        part[tid] = s;
        __syncthreads();
        if (lane == 0) {
            int base = j << 2;
            float v = b2[j] + part[base] + part[base + 1] + part[base + 2] + part[base + 3];
            h2[j] = silu_f(v);
        }
        __syncthreads();

        // Output linear: sixteen threads per output coordinate.
        int oj = tid >> 4;
        int olane = tid & 15;
        s = 0.0f;
        #pragma unroll 4
        for (int i = olane; i < H; i += 16) {
            s = fmaf(h2[i], w3[oj * H + i], s);
        }
        part[tid] = s;
        __syncthreads();
        if (olane == 0) {
            int base = oj << 4;
            float v = b3[oj];
            #pragma unroll
            for (int q = 0; q < 16; ++q) v += part[base + q];
            x[oj] = fmaf(inv_steps, v, x[oj]);
        }
        __syncthreads();
    }

    if (tid < D) out[b * D + tid] = x[tid];
}

torch::Tensor flow_forward(torch::Tensor x0,
                           torch::Tensor w1, torch::Tensor b1,
                           torch::Tensor w2, torch::Tensor b2,
                           torch::Tensor w3, torch::Tensor b3,
                           int64_t num_steps) {
    CHECK_CUDA(x0);
    CHECK_CONTIGUOUS(x0);
    CHECK_CONTIGUOUS(w1); CHECK_CONTIGUOUS(b1);
    CHECK_CONTIGUOUS(w2); CHECK_CONTIGUOUS(b2);
    CHECK_CONTIGUOUS(w3); CHECK_CONTIGUOUS(b3);
    TORCH_CHECK(x0.dtype() == torch::kFloat32 && x0.dim() == 2 && x0.size(1) == 64,
                "expected a contiguous CUDA fp32 tensor of shape (B, 64)");

    int B = (int)x0.size(0);
    auto out = torch::empty_like(x0);
    flow_kernel_64_4way<<<B, 1024, 0, at::cuda::getCurrentCUDAStream()>>>(
        x0.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>(),
        out.data_ptr<float>(), B, (int)num_steps);
    return out;
}
"""

_ext = load_inline(
    name="flow_fused_ext_final",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["flow_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    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.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:
        return _ext.flow_forward(
            x0.contiguous(),
            self.velocity[0].weight.contiguous(), self.velocity[0].bias.contiguous(),
            self.velocity[2].weight.contiguous(), self.velocity[2].bias.contiguous(),
            self.velocity[4].weight.contiguous(), self.velocity[4].bias.contiguous(),
            self.num_steps,
        )