import torch
import torch.nn as nn
import torch.nn.functional as F

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


def _linear_beta_schedule(num_timesteps: int, beta_start: float, beta_end: float) -> torch.Tensor:
    return torch.linspace(beta_start, beta_end, num_timesteps)


class Model(nn.Module):
    """
    Fused **multi-step** diffusion denoising (VP-DDPM-style) in a single ``forward``.

    Instead of one scheduler step per Python call, this reference unrolls
    ``num_fused_steps`` consecutive **deterministic** reverse updates (posterior mean,
    no sampler noise)—a typical fusion target for robotics / world-model diffusion
    heads where latency matters.

    Each inner step: predict ``ε_θ(x, t)`` with a small MLP, then apply the standard
    DDPM mean update from discrete time ``t`` to ``t-1``. Batch elements may start
    from different ``t_start``; inputs must satisfy ``t_start >= num_fused_steps`` so
    indices stay valid.
    """

    def __init__(self, dim: int, num_fused_steps: int, total_timesteps: int = 1000):
        super().__init__()
        assert num_fused_steps >= 1
        assert total_timesteps >= num_fused_steps + 1
        self.dim = dim
        self.num_fused_steps = num_fused_steps
        self.total_timesteps = total_timesteps

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

        betas = _linear_beta_schedule(total_timesteps, 1e-4, 0.02)
        alphas = 1.0 - betas
        alphas_cumprod = torch.cumprod(alphas, dim=0)
        self.register_buffer("betas", betas)
        self.register_buffer("alphas", alphas)
        self.register_buffer("alphas_cumprod", alphas_cumprod)
        # x_{t-1} = (1/sqrt(alpha_t)) * (x_t - (beta_t / sqrt(1-alpha_bar_t)) * eps)
        self.register_buffer("sqrt_recip_alphas", torch.sqrt(1.0 / alphas))
        self.register_buffer(
            "coef_eps",
            betas / torch.sqrt(1.0 - alphas_cumprod),
        )

    def forward(self, x: torch.Tensor, t_start: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, dim) noisy latent at discrete time ``t_start`` (same semantics as DDPM).
            t_start: (B,) int64 timesteps in ``[num_fused_steps, total_timesteps - 1]``.
        Returns:
            (B, dim) latent after ``num_fused_steps`` fused denoise steps toward 0.
        """
        B, d = x.shape
        assert d == self.dim
        assert t_start.shape == (B,)
        # Harness paths (measure_ref_program_time / eval) cast activations to fp32;
        # discrete timesteps must survive as whole-valued floats → long indices.
        t_start = t_start.long()

        x_cur = x
        for k in range(self.num_fused_steps):
            t_idx = t_start - k
            # ``get_inputs`` samples t_start >= num_fused_steps so t_idx >= 1 always.
            t_norm = t_idx.float() / float(self.total_timesteps)
            eps = self.eps_net(torch.cat([x_cur, t_norm.unsqueeze(-1)], dim=-1))

            c1 = self.sqrt_recip_alphas[t_idx].view(B, 1)
            c2 = self.coef_eps[t_idx].view(B, 1)
            x_cur = c1 * (x_cur - c2 * eps)

        return x_cur


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

    def __init__(self, dim: int, num_fused_steps: int, total_timesteps: int = 1000):
        super().__init__()
        self._impl = Model(dim, num_fused_steps, total_timesteps=total_timesteps)

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


batch_size = 8
dim = 64
num_fused_steps = 16
total_timesteps = 1000


def get_inputs():
    p = popcorn_pri
    x = torch.randn(p.jitter_int(batch_size), dim)
    low = num_fused_steps
    high = total_timesteps
    t_start = torch.randint(low, high, (p.jitter_int(batch_size),), dtype=torch.long)
    return [x, t_start]


def get_init_inputs():
    return [dim, num_fused_steps, total_timesteps]

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline
import math

def _linear_beta_schedule(num_timesteps: int, beta_start: float, beta_end: float) -> torch.Tensor:
    return torch.linspace(beta_start, beta_end, num_timesteps)

# CUDA kernel code - fully fused multi-step diffusion denoising
cuda_source = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

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

__global__ void __launch_bounds__(256) fused_diffusion_kernel(
    const float* __restrict__ x,
    const int64_t* __restrict__ t_start,
    const float* __restrict__ w1,
    const float* __restrict__ b1,
    const float* __restrict__ w2,
    const float* __restrict__ b2,
    const float* __restrict__ w3,
    const float* __restrict__ b3,
    const float* __restrict__ sqrt_recip_alphas,
    const float* __restrict__ coef_eps,
    float* __restrict__ out,
    int batch_size, int dim, int num_fused_steps, int total_timesteps
) {
    const int dim4 = dim * 4;
    const int dim_plus1 = dim + 1;
    
    const int batch_idx = blockIdx.x;
    if (batch_idx >= batch_size) return;
    
    const int tid = threadIdx.x;
    const int block_size = blockDim.x;
    
    const int t_start_val = (int)t_start[batch_idx];
    
    __shared__ float x_cur[256];
    __shared__ float hidden1[1024];
    __shared__ float hidden2[1024];
    __shared__ float eps[256];
    
    for (int i = tid; i < dim; i += block_size) {
        x_cur[i] = x[batch_idx * dim + i];
    }
    __syncthreads();
    
    for (int k = 0; k < num_fused_steps; k++) {
        int t_idx = t_start_val - k;
        float t_norm = (float)t_idx / (float)total_timesteps;
        
        for (int i = tid; i < dim4; i += block_size) {
            float sum = b1[i];
            #pragma unroll 8
            for (int j = 0; j < dim; j++) {
                sum += w1[i * dim_plus1 + j] * x_cur[j];
            }
            sum += w1[i * dim_plus1 + dim] * t_norm;
            hidden1[i] = sum;
        }
        __syncthreads();
        
        for (int i = tid; i < dim4; i += block_size) {
            hidden1[i] = silu(hidden1[i]);
        }
        __syncthreads();
        
        for (int i = tid; i < dim4; i += block_size) {
            float sum = b2[i];
            #pragma unroll 8
            for (int j = 0; j < dim4; j++) {
                sum += w2[i * dim4 + j] * hidden1[j];
            }
            hidden2[i] = sum;
        }
        __syncthreads();
        
        for (int i = tid; i < dim4; i += block_size) {
            hidden2[i] = silu(hidden2[i]);
        }
        __syncthreads();
        
        for (int i = tid; i < dim; i += block_size) {
            float sum = b3[i];
            #pragma unroll 8
            for (int j = 0; j < dim4; j++) {
                sum += w3[i * dim4 + j] * hidden2[j];
            }
            eps[i] = sum;
        }
        __syncthreads();
        
        float c1 = sqrt_recip_alphas[t_idx];
        float c2 = coef_eps[t_idx];
        
        for (int i = tid; i < dim; i += block_size) {
            x_cur[i] = c1 * (x_cur[i] - c2 * eps[i]);
        }
        __syncthreads();
    }
    
    for (int i = tid; i < dim; i += block_size) {
        out[batch_idx * dim + i] = x_cur[i];
    }
}

torch::Tensor fused_diffusion_forward(
    torch::Tensor x,
    torch::Tensor t_start,
    torch::Tensor w1,
    torch::Tensor b1,
    torch::Tensor w2,
    torch::Tensor b2,
    torch::Tensor w3,
    torch::Tensor b3,
    torch::Tensor sqrt_recip_alphas,
    torch::Tensor coef_eps,
    int num_fused_steps,
    int total_timesteps
) {
    const int batch_size = x.size(0);
    const int dim = x.size(1);
    
    auto out = torch::empty_like(x);
    
    fused_diffusion_kernel<<<batch_size, 256>>>(
        x.data_ptr<float>(),
        t_start.data_ptr<int64_t>(),
        w1.data_ptr<float>(),
        b1.data_ptr<float>(),
        w2.data_ptr<float>(),
        b2.data_ptr<float>(),
        w3.data_ptr<float>(),
        b3.data_ptr<float>(),
        sqrt_recip_alphas.data_ptr<float>(),
        coef_eps.data_ptr<float>(),
        out.data_ptr<float>(),
        batch_size, dim, num_fused_steps, total_timesteps
    );
    
    return out;
}
"""

cpp_source = r"""
torch::Tensor fused_diffusion_forward(
    torch::Tensor x,
    torch::Tensor t_start,
    torch::Tensor w1,
    torch::Tensor b1,
    torch::Tensor w2,
    torch::Tensor b2,
    torch::Tensor w3,
    torch::Tensor b3,
    torch::Tensor sqrt_recip_alphas,
    torch::Tensor coef_eps,
    int num_fused_steps,
    int total_timesteps
);
"""

fused_extension = load_inline(
    name='fused_diffusion_ext',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['fused_diffusion_forward'],
    extra_cuda_cflags=['-O3', '-use_fast_math'],
    verbose=False
)


class ModelNew(nn.Module):
    def __init__(self, dim: int, num_fused_steps: int, total_timesteps: int = 1000):
        super().__init__()
        assert num_fused_steps >= 1
        assert total_timesteps >= num_fused_steps + 1
        self.dim = dim
        self.num_fused_steps = num_fused_steps
        self.total_timesteps = total_timesteps

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

        betas = _linear_beta_schedule(total_timesteps, 1e-4, 0.02)
        alphas = 1.0 - betas
        alphas_cumprod = torch.cumprod(alphas, dim=0)
        self.register_buffer("betas", betas)
        self.register_buffer("alphas", alphas)
        self.register_buffer("alphas_cumprod", alphas_cumprod)
        self.register_buffer("sqrt_recip_alphas", torch.sqrt(1.0 / alphas))
        self.register_buffer(
            "coef_eps",
            betas / torch.sqrt(1.0 - alphas_cumprod),
        )

    def forward(self, x: torch.Tensor, t_start: torch.Tensor) -> torch.Tensor:
        B, d = x.shape
        assert d == self.dim
        assert t_start.shape == (B,)
        t_start = t_start.long()
        
        w1 = self.eps_net[0].weight
        b1 = self.eps_net[0].bias
        w2 = self.eps_net[2].weight
        b2 = self.eps_net[2].bias
        w3 = self.eps_net[4].weight
        b3 = self.eps_net[4].bias
        
        return fused_extension.fused_diffusion_forward(
            x.contiguous(), t_start.contiguous(),
            w1.contiguous(), b1.contiguous(),
            w2.contiguous(), b2.contiguous(),
            w3.contiguous(), b3.contiguous(),
            self.sqrt_recip_alphas.contiguous(),
            self.coef_eps.contiguous(),
            self.num_fused_steps, self.total_timesteps
        )


batch_size = 8
dim = 64
num_fused_steps = 16
total_timesteps = 1000


def get_inputs():
    x = torch.randn(batch_size, dim)
    low = num_fused_steps
    high = total_timesteps
    t_start = torch.randint(low, high, (batch_size,), dtype=torch.long)
    return [x, t_start]


def get_init_inputs():
    return [dim, num_fused_steps, total_timesteps]