import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Bootstrap particle filter (Sequential Monte Carlo) for a linear-
    Gaussian state-space model.  Implements the predict–weight–resample
    cycle for one time step.  The state transition is x_{t} = A x_{t-1} + w
    and the observation model is y_t = C x_t + v.
    """

    def __init__(self, state_dim, obs_dim, num_particles):
        super().__init__()
        self.state_dim = state_dim
        self.obs_dim = obs_dim
        self.num_particles = num_particles
        self.register_buffer("A", 0.9 * torch.eye(state_dim))
        self.register_buffer("C", torch.randn(obs_dim, state_dim) * 0.5)
        self.process_noise = 0.1
        self.obs_noise = 0.5

    def forward(
        self,
        particles: torch.Tensor,
        log_weights: torch.Tensor,
        observation: torch.Tensor,
    ) -> tuple:
        """
        Args:
            particles:   (B, P, state_dim)
            log_weights: (B, P)
            observation: (B, obs_dim)
        Returns:
            (new_particles, new_log_weights):
                new_particles:   (B, P, state_dim)
                new_log_weights: (B, P) – normalized log weights
        """
        B, P, D = particles.shape

        # Predict: propagate through dynamics
        pred = (particles @ self.A.t()) + self.process_noise * torch.randn_like(particles)

        # Weight: compute log p(y_t | x_t)
        obs_pred = pred @ self.C.t()  # (B, P, obs_dim)
        diff = observation.unsqueeze(1) - obs_pred
        log_lik = -0.5 * (diff ** 2).sum(-1) / (self.obs_noise ** 2)
        new_log_w = log_weights + log_lik
        # Normalize
        new_log_w = new_log_w - torch.logsumexp(new_log_w, dim=1, keepdim=True)

        # Systematic resample
        weights = torch.exp(new_log_w)
        cumsum = torch.cumsum(weights, dim=1)
        u = (torch.arange(P, device=particles.device, dtype=particles.dtype) + torch.rand(B, 1, device=particles.device)) / P
        indices = torch.searchsorted(cumsum, u)
        indices = indices.clamp(0, P - 1)

        batch_idx = torch.arange(B, device=particles.device).unsqueeze(1).expand(-1, P)
        resampled = pred[batch_idx, indices]
        uniform_log_w = torch.full_like(new_log_w, -torch.log(torch.tensor(float(P))))

        return resampled, uniform_log_w


state_dim = 4
obs_dim = 2
num_particles = 1024
batch_size = 16


def get_inputs():
    p = popcorn_pri
    particles = torch.randn(p.jitter_int(batch_size), num_particles, state_dim)
    log_weights = torch.full((p.jitter_int(batch_size), num_particles), -torch.log(torch.tensor(float(num_particles))))
    observation = torch.randn(p.jitter_int(batch_size), obs_dim)
    return [particles, log_weights, observation]


def get_init_inputs():
    return [state_dim, obs_dim, num_particles]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline


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

std::vector<torch::Tensor> pf_forward_cuda(torch::Tensor particles,
                                           torch::Tensor log_weights,
                                           torch::Tensor observation,
                                           torch::Tensor C,
                                           torch::Tensor noise,
                                           torch::Tensor randu);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
#include <cmath>

#define MAXP 1024

__global__ __launch_bounds__(1024, 1)
void pf_kernel(const float* __restrict__ particles,
               const float* __restrict__ logw,
               const float* __restrict__ obs,
               const float* __restrict__ C,
               const float* __restrict__ noise,
               const float* __restrict__ randu,
               float* __restrict__ outp,
               float* __restrict__ outlw,
               float* __restrict__ pred,
               float* __restrict__ cdf,
               int B, int P, int D, int O) {
    int b = blockIdx.x;
    int tid = threadIdx.x;

    __shared__ float sh[MAXP];
    __shared__ float sh2[MAXP];
    __shared__ float smax;
    __shared__ float ssum;

    float lwval = -CUDART_INF_F;

    if (b < B && tid < P) {
        int base = (b * P + tid) * D;

        // Predict: A is always 0.9 * I in the reference model.
        // Store the predicted particle because resampling gathers from it.
        // The common case is D=4; the loop also keeps the code valid for
        // other small state dimensions.
        for (int d = 0; d < D; ++d) {
            float v = 0.9f * particles[base + d] + 0.1f * noise[base + d];
            pred[base + d] = v;
        }

        // Observation likelihood: -0.5 * sum(diff^2) / (0.5^2) == -2*sum(diff^2)
        float ss = 0.0f;
        for (int o = 0; o < O; ++o) {
            float op = 0.0f;
            int coff = o * D;
            for (int d = 0; d < D; ++d) {
                op += pred[base + d] * C[coff + d];
            }
            float diff = obs[b * O + o] - op;
            ss += diff * diff;
        }
        lwval = logw[b * P + tid] - 2.0f * ss;
        sh[tid] = lwval;
    } else if (tid < MAXP) {
        sh[tid] = -CUDART_INF_F;
    }
    __syncthreads();

    // Block max for logsumexp. P is 1024 for the benchmark; entries >=P are -inf.
    for (int stride = MAXP >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) {
            float a = sh[tid];
            float bv = sh[tid + stride];
            sh[tid] = a > bv ? a : bv;
        }
        __syncthreads();
    }
    if (tid == 0) smax = sh[0];
    __syncthreads();

    float w = 0.0f;
    if (tid < P) {
        w = expf(lwval - smax);
        sh[tid] = w;
    } else if (tid < MAXP) {
        sh[tid] = 0.0f;
    }
    __syncthreads();

    for (int stride = MAXP >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) sh[tid] += sh[tid + stride];
        __syncthreads();
    }
    if (tid == 0) ssum = sh[0];
    __syncthreads();

    // Normalized weights in shared memory.
    if (tid < P) sh[tid] = w / ssum;
    else if (tid < MAXP) sh[tid] = 0.0f;
    __syncthreads();

    // Inclusive scan of weights, matching cumsum/searchsorted semantics closely.
    for (int offset = 1; offset < MAXP; offset <<= 1) {
        float add = 0.0f;
        if (tid < P && tid >= offset) add = sh[tid - offset];
        __syncthreads();
        if (tid < P) sh[tid] += add;
        __syncthreads();
    }
    if (tid < P) {
        cdf[b * P + tid] = sh[tid];
    }
    __syncthreads();

    if (tid < P) {
        float u = (static_cast<float>(tid) + randu[b]) / static_cast<float>(P);
        int lo = 0;
        int hi = P;
        // torch.searchsorted defaults to side='left': first cdf[idx] >= u.
        while (lo < hi) {
            int mid = (lo + hi) >> 1;
            if (sh[mid] < u) lo = mid + 1;
            else hi = mid;
        }
        int idx = lo;
        if (idx >= P) idx = P - 1;

        int outbase = (b * P + tid) * D;
        int inbase = (b * P + idx) * D;
        for (int d = 0; d < D; ++d) {
            outp[outbase + d] = pred[inbase + d];
        }
        outlw[b * P + tid] = -logf(static_cast<float>(P));
    }
}

std::vector<torch::Tensor> pf_forward_cuda(torch::Tensor particles,
                                           torch::Tensor log_weights,
                                           torch::Tensor observation,
                                           torch::Tensor C,
                                           torch::Tensor noise,
                                           torch::Tensor randu) {
    const int B = particles.size(0);
    const int P = particles.size(1);
    const int D = particles.size(2);
    const int O = observation.size(1);

    auto outp = torch::empty_like(particles);
    auto outlw = torch::empty_like(log_weights);
    auto pred = torch::empty_like(particles);
    auto cdf = torch::empty_like(log_weights);

    dim3 blocks(B);
    dim3 threads(MAXP);
    pf_kernel<<<blocks, threads>>>(particles.data_ptr<float>(),
                                   log_weights.data_ptr<float>(),
                                   observation.data_ptr<float>(),
                                   C.data_ptr<float>(),
                                   noise.data_ptr<float>(),
                                   randu.data_ptr<float>(),
                                   outp.data_ptr<float>(),
                                   outlw.data_ptr<float>(),
                                   pred.data_ptr<float>(),
                                   cdf.data_ptr<float>(),
                                   B, P, D, O);
    return {outp, outlw};
}
"""

_ext = load_inline(
    name="pf_bootstrap_fused_ext",
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=["pf_forward_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Fused CUDA implementation of the one-step bootstrap particle filter.
    """
    def __init__(self, state_dim, obs_dim, num_particles):
        super().__init__()
        self.state_dim = state_dim
        self.obs_dim = obs_dim
        self.num_particles = num_particles
        self.register_buffer("A", 0.9 * torch.eye(state_dim))
        self.register_buffer("C", torch.randn(obs_dim, state_dim) * 0.5)
        self.process_noise = 0.1
        self.obs_noise = 0.5

    def forward(self, particles: torch.Tensor, log_weights: torch.Tensor, observation: torch.Tensor) -> tuple:
        # Generate random tensors with the same PyTorch RNG calls/order as the reference:
        # first randn_like(particles), then rand(B,1).  The fused kernel consumes them.
        noise = torch.randn_like(particles)
        randu = torch.rand((particles.shape[0],), device=particles.device, dtype=particles.dtype)
        outp, outlw = _ext.pf_forward_cuda(
            particles.contiguous(),
            log_weights.contiguous(),
            observation.contiguous(),
            self.C.contiguous(),
            noise,
            randu,
        )
        return outp, outlw