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

# Fused probability-flow ODE integration kernel
cuda_source = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

// SiLU activation: x * sigmoid(x)
__device__ __forceinline__ float silu(float x) {
    return x * (1.0f / (1.0f + expf(-x)));
}

// Fused MLP forward pass for velocity network
template <int IN_DIM, int HIDDEN_DIM, int OUT_DIM, int BLOCK_SIZE>
__device__ void fused_mlp_forward(
    const float* __restrict__ input,      // [IN_DIM]
    float* __restrict__ output,           // [OUT_DIM]
    const float* __restrict__ w1,         // [HIDDEN_DIM, IN_DIM]
    const float* __restrict__ b1,         // [HIDDEN_DIM]
    const float* __restrict__ w2,         // [HIDDEN_DIM, HIDDEN_DIM]
    const float* __restrict__ b2,         // [HIDDEN_DIM]
    const float* __restrict__ w3,         // [OUT_DIM, HIDDEN_DIM]
    const float* __restrict__ b3,         // [OUT_DIM]
    float* __restrict__ shared_buf        // shared memory buffer
) {
    // Stage 1: Linear1 + SiLU
    float* hidden1 = shared_buf;  // [HIDDEN_DIM]
    
    for (int h = threadIdx.x; h < HIDDEN_DIM; h += BLOCK_SIZE) {
        float sum = b1[h];
        for (int i = 0; i < IN_DIM; i++) {
            sum += w1[h * IN_DIM + i] * input[i];
        }
        hidden1[h] = silu(sum);
    }
    __syncthreads();
    
    // Stage 2: Linear2 + SiLU
    float* hidden2 = shared_buf + HIDDEN_DIM;  // [HIDDEN_DIM]
    
    for (int h = threadIdx.x; h < HIDDEN_DIM; h += BLOCK_SIZE) {
        float sum = b2[h];
        for (int i = 0; i < HIDDEN_DIM; i++) {
            sum += w2[h * HIDDEN_DIM + i] * hidden1[i];
        }
        hidden2[h] = silu(sum);
    }
    __syncthreads();
    
    // Stage 3: Linear3 (output layer)
    for (int o = threadIdx.x; o < OUT_DIM; o += BLOCK_SIZE) {
        float sum = b3[o];
        for (int i = 0; i < HIDDEN_DIM; i++) {
            sum += w3[o * HIDDEN_DIM + i] * hidden2[i];
        }
        output[o] = sum;
    }
    __syncthreads();
}

// Main kernel: fused Euler integration
template <int DIM, int HIDDEN_DIM, int NUM_STEPS, int BLOCK_SIZE>
__global__ void fused_euler_ode_kernel(
    const float* __restrict__ x0,          // [B, DIM]
    float* __restrict__ x_out,             // [B, DIM]
    const float* __restrict__ w1,          // [HIDDEN_DIM, DIM+1]
    const float* __restrict__ b1,          // [HIDDEN_DIM]
    const float* __restrict__ w2,          // [HIDDEN_DIM, HIDDEN_DIM]
    const float* __restrict__ b2,          // [HIDDEN_DIM]
    const float* __restrict__ w3,          // [DIM, HIDDEN_DIM]
    const float* __restrict__ b3,          // [DIM]
    const float dt                         // time step = 1.0 / NUM_STEPS
) {
    const int batch_idx = blockIdx.x;
    const int tid = threadIdx.x;
    
    // Shared memory: input buffer + hidden1 + hidden2 + velocity
    extern __shared__ float shared_mem[];
    float* input_buf = shared_mem;              // [DIM+1]
    float* hidden_buf = shared_mem + (DIM + 1); // [2 * HIDDEN_DIM]
    float* velocity_buf = hidden_buf + 2 * HIDDEN_DIM;  // [DIM]
    
    // Load initial x into registers
    float x_reg[DIM];
    const float* x0_ptr = x0 + batch_idx * DIM;
    
    #pragma unroll
    for (int d = tid; d < DIM; d += BLOCK_SIZE) {
        x_reg[d] = x0_ptr[d];
    }
    __syncthreads();
    
    // Euler integration loop
    for (int step = 0; step < NUM_STEPS; step++) {
        // Compute time t = (k + 0.5) / num_steps
        const float t = (float(step) + 0.5f) / float(NUM_STEPS);
        
        // Prepare input: [x, t]
        for (int d = tid; d < DIM; d += BLOCK_SIZE) {
            input_buf[d] = x_reg[d];
        }
        if (tid == 0) {
            input_buf[DIM] = t;
        }
        __syncthreads();
        
        // Compute velocity through fused MLP
        fused_mlp_forward<DIM + 1, HIDDEN_DIM, DIM, BLOCK_SIZE>(
            input_buf, velocity_buf,
            w1, b1, w2, b2, w3, b3,
            hidden_buf
        );
        __syncthreads();
        
        // Euler update: x = x + dt * v
        for (int d = tid; d < DIM; d += BLOCK_SIZE) {
            x_reg[d] += dt * velocity_buf[d];
        }
        __syncthreads();
    }
    
    // Write output
    float* x_out_ptr = x_out + batch_idx * DIM;
    for (int d = tid; d < DIM; d += BLOCK_SIZE) {
        x_out_ptr[d] = x_reg[d];
    }
}

// Host function to launch kernel
void launch_fused_euler_ode(
    const torch::Tensor& x0,
    torch::Tensor& x_out,
    const torch::Tensor& w1,
    const torch::Tensor& b1,
    const torch::Tensor& w2,
    const torch::Tensor& b2,
    const torch::Tensor& w3,
    const torch::Tensor& b3,
    int dim, int hidden_dim, int num_steps
) {
    const int batch_size = x0.size(0);
    const float dt = 1.0f / float(num_steps);
    
    const int block_size = 128;  // Tunable
    
    // Shared memory size: input(DIM+1) + 2*hidden + output(DIM)
    size_t shared_mem_size = (dim + 1 + 2 * hidden_dim + dim) * sizeof(float);
    
    // Launch configuration
    dim3 grid(batch_size);
    dim3 block(block_size);
    
    fused_euler_ode_kernel<64, 256, 32, 128><<<grid, block, shared_mem_size>>>(
        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>(),
        dt
    );
}
"""

cpp_source = r"""
void launch_fused_euler_ode(
    const torch::Tensor& x0,
    torch::Tensor& x_out,
    const torch::Tensor& w1,
    const torch::Tensor& b1,
    const torch::Tensor& w2,
    const torch::Tensor& b2,
    const torch::Tensor& w3,
    const torch::Tensor& b3,
    int dim, int hidden_dim, int num_steps
);
"""

# Load the CUDA extension
fused_ode_module = load_inline(
    name='fused_ode_extension',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['launch_fused_euler_ode'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    verbose=True
)


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.hidden_dim = dim * 4
        
        # Create the velocity network (same architecture as reference)
        self.velocity = nn.Sequential(
            nn.Linear(dim + 1, self.hidden_dim),
            nn.SiLU(),
            nn.Linear(self.hidden_dim, self.hidden_dim),
            nn.SiLU(),
            nn.Linear(self.hidden_dim, dim),
        )

    def forward(self, x0: torch.Tensor) -> torch.Tensor:
        B, d = x0.shape
        assert d == self.dim
        
        # Ensure contiguous float32 tensors
        x0_f32 = x0.float().contiguous()
        
        # Get weight tensors from velocity network
        # Sequential: [Linear, SiLU, Linear, SiLU, Linear]
        # Indices:       0       1      2       3       4
        w1 = self.velocity[0].weight.data.float().contiguous()  # [hidden, dim+1]
        b1 = self.velocity[0].bias.data.float().contiguous()    # [hidden]
        w2 = self.velocity[2].weight.data.float().contiguous()  # [hidden, hidden]
        b2 = self.velocity[2].bias.data.float().contiguous()    # [hidden]
        w3 = self.velocity[4].weight.data.float().contiguous()  # [dim, hidden]
        b3 = self.velocity[4].bias.data.float().contiguous()    # [dim]
        
        # Allocate output
        x_out = torch.empty_like(x0_f32)
        
        # Launch fused kernel
        fused_ode_module.launch_fused_euler_ode(
            x0_f32, x_out,
            w1, b1, w2, b2, w3, b3,
            self.dim, self.hidden_dim, self.num_steps
        )
        
        return x_out.to(x0.dtype)