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


class Model(nn.Module):
    """
    Stack of planar normalizing flow layers.  Each layer applies the
    invertible transform f(z) = z + u * tanh(w^T z + b), with a
    log-determinant-Jacobian correction.  Used to transform a simple
    base distribution (e.g. standard Gaussian) into a complex posterior.
    """

    def __init__(self, dim, num_flows):
        super().__init__()
        self.dim = dim
        self.num_flows = num_flows
        self.w = nn.ParameterList([nn.Parameter(torch.randn(dim)) for _ in range(num_flows)])
        self.u = nn.ParameterList([nn.Parameter(torch.randn(dim)) for _ in range(num_flows)])
        self.b = nn.ParameterList([nn.Parameter(torch.zeros(1)) for _ in range(num_flows)])

    def _constrain_u(self, w, u):
        """Ensure invertibility: w^T u >= -1."""
        wtu = (w * u).sum()
        m = -1 + F.softplus(wtu)
        u_hat = u + (m - wtu) * w / (w ** 2).sum()
        return u_hat

    def forward(self, z: torch.Tensor) -> tuple:
        """
        Args:
            z: (B, dim) – samples from base distribution
        Returns:
            (z_out, sum_log_det):
                z_out:       (B, dim) – transformed samples
                sum_log_det: (B,)     – total log |det df/dz|
        """
        sum_log_det = torch.zeros(z.shape[0], device=z.device)

        for i in range(self.num_flows):
            w = self.w[i]
            u = self._constrain_u(w, self.u[i])
            b = self.b[i]

            linear = z @ w + b  # (B,)
            h = torch.tanh(linear)
            h_prime = 1.0 - h ** 2

            z = z + u.unsqueeze(0) * h.unsqueeze(-1)
            psi = h_prime.unsqueeze(-1) * w.unsqueeze(0)  # (B, dim)
            log_det = torch.log(torch.abs(1.0 + psi @ u) + 1e-8)
            sum_log_det = sum_log_det + log_det

        return z, sum_log_det


dim = 32
num_flows = 8
batch_size = 1024


def get_inputs():
    return [torch.randn(popcorn_pri.jitter_int(batch_size), dim)]


def get_init_inputs():
    return [dim, num_flows]

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

_cpp = r'''
#include <torch/extension.h>
#include <vector>
std::vector<torch::Tensor> planar_forward(torch::Tensor z, std::vector<torch::Tensor> w_list, std::vector<torch::Tensor> u_list, std::vector<torch::Tensor> b_list);
'''

_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
#include <cmath>
#include <c10/cuda/CUDAGuard.h>
#include <ATen/cuda/CUDAContext.h>

static __device__ __forceinline__ float warp_sum(float v) {
    unsigned mask = 0xffffffffu;
    v += __shfl_down_sync(mask, v, 16);
    v += __shfl_down_sync(mask, v, 8);
    v += __shfl_down_sync(mask, v, 4);
    v += __shfl_down_sync(mask, v, 2);
    v += __shfl_down_sync(mask, v, 1);
    return v;
}
static __device__ __forceinline__ float softplus_pt(float x) {
    if (x > 20.0f) return x;
    return log1pf(expf(x));
}
static __device__ __forceinline__ float tf32_rna(float x) {
#if __CUDA_ARCH__ >= 800
    float y;
    asm("cvt.rna.tf32.f32 %0, %1;" : "=f"(y) : "f"(x));
    return y;
#else
    return x;
#endif
}

__global__ __launch_bounds__(32, 16)
void planar_kernel(const float* __restrict__ z_in, float* __restrict__ z_out, float* __restrict__ ld_out,
                   const float* const* __restrict__ w_ptrs, const float* const* __restrict__ u_ptrs, const float* const* __restrict__ b_ptrs,
                   int B, int D, int F) {
    int row = blockIdx.x;
    int lane = threadIdx.x;
    if (lane >= D) return;
    float z = z_in[row * D + lane];
    float sum_ld = 0.0f;
    #pragma unroll
    for (int f = 0; f < 8; ++f) {
        if (f >= F) break;
        const float* __restrict__ wv = w_ptrs[f];
        const float* __restrict__ uv = u_ptrs[f];
        float w = wv[lane];
        float u = uv[lane];
        float wtu = warp_sum(w * u);
        float w2 = warp_sum(w * w);
        // torch matmul on Ampere/Hopper normally uses TF32 for fp32 GEMV/GEMM; emulate it for the two matvecs.
        float lin = warp_sum(tf32_rna(z) * tf32_rna(w)) + b_ptrs[f][0];
        wtu = __shfl_sync(0xffffffffu, wtu, 0);
        w2 = __shfl_sync(0xffffffffu, w2, 0);
        lin = __shfl_sync(0xffffffffu, lin, 0);
        float m = -1.0f + softplus_pt(wtu);
        float coeff = (m - wtu) / w2;
        float uhat = u + coeff * w;
        float h = tanhf(lin);
        float hp = 1.0f - h * h;
        z += uhat * h;
        float wu = warp_sum(tf32_rna(hp * w) * tf32_rna(uhat));
        wu = __shfl_sync(0xffffffffu, wu, 0);
        if (lane == 0) sum_ld += logf(fabsf(1.0f + wu) + 1.0e-8f);
    }
    z_out[row * D + lane] = z;
    if (lane == 0) ld_out[row] = sum_ld;
}

std::vector<torch::Tensor> planar_forward(torch::Tensor z,
                                          std::vector<torch::Tensor> w_list,
                                          std::vector<torch::Tensor> u_list,
                                          std::vector<torch::Tensor> b_list) {
    const at::cuda::OptionalCUDAGuard device_guard(device_of(z));
    int64_t B64 = z.size(0), D64 = z.size(1);
    int F = (int)w_list.size();
    auto out = torch::empty_like(z);
    auto ld = torch::empty({B64}, z.options());
    std::vector<const float*> hw(F), hu(F), hb(F);
    for (int i = 0; i < F; ++i) { hw[i]=w_list[i].data_ptr<float>(); hu[i]=u_list[i].data_ptr<float>(); hb[i]=b_list[i].data_ptr<float>(); }
    auto opts = torch::TensorOptions().device(z.device()).dtype(torch::kInt64);
    auto wp = torch::empty({F}, opts), up = torch::empty({F}, opts), bp = torch::empty({F}, opts);
    auto stream = at::cuda::getCurrentCUDAStream();
    cudaMemcpyAsync(wp.data_ptr<int64_t>(), hw.data(), F*sizeof(const float*), cudaMemcpyHostToDevice, stream);
    cudaMemcpyAsync(up.data_ptr<int64_t>(), hu.data(), F*sizeof(const float*), cudaMemcpyHostToDevice, stream);
    cudaMemcpyAsync(bp.data_ptr<int64_t>(), hb.data(), F*sizeof(const float*), cudaMemcpyHostToDevice, stream);
    planar_kernel<<<(int)B64, 32, 0, stream>>>(z.data_ptr<float>(), out.data_ptr<float>(), ld.data_ptr<float>(),
        reinterpret_cast<const float* const*>(wp.data_ptr<int64_t>()), reinterpret_cast<const float* const*>(up.data_ptr<int64_t>()), reinterpret_cast<const float* const*>(bp.data_ptr<int64_t>()),
        (int)B64, (int)D64, F);
    return {out, ld};
}
'''

_ext = load_inline(name='planar_flow_final_ext', cpp_sources=_cpp, cuda_sources=_src, functions=['planar_forward'], verbose=False, extra_cuda_cflags=['-O3'], extra_cflags=['-O3'])

class ModelNew(nn.Module):
    def __init__(self, dim, num_flows):
        super().__init__()
        self.dim = dim
        self.num_flows = num_flows
        self.w = nn.ParameterList([nn.Parameter(torch.randn(dim)) for _ in range(num_flows)])
        self.u = nn.ParameterList([nn.Parameter(torch.randn(dim)) for _ in range(num_flows)])
        self.b = nn.ParameterList([nn.Parameter(torch.zeros(1)) for _ in range(num_flows)])
    def forward(self, z: torch.Tensor) -> tuple:
        return tuple(_ext.planar_forward(z, list(self.w), list(self.u), list(self.b)))