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

cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <vector>
#include <cmath>

__device__ float softplus(float x) {
    return log1pf(expf(x));
}

__global__ void planar_flow_kernel(
    const float* __restrict__ z_in,
    float* __restrict__ z_out,
    float* __restrict__ sum_log_det,
    const float* __restrict__ w,
    const float* __restrict__ u,
    const float* __restrict__ b,
    int batch_size,
    int dim
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= batch_size) return;

    double linear = b[0];
    for (int d = 0; d < dim; d++) {
        linear += (double)z_in[idx * dim + d] * (double)w[d];
    }
    float h = tanhf((float)linear);
    float h_prime = 1.0f - h * h;

    double w_norm2 = 0.0;
    double wtu = 0.0;
    for (int d = 0; d < dim; d++) {
        w_norm2 += (double)w[d] * (double)w[d];
        wtu += (double)w[d] * (double)u[d];
    }
    float m = -1.0f + softplus((float)wtu);
    float scale = ((float)m - (float)wtu) / (float)w_norm2;

    double psi_dot_u = 0.0;
    for (int d = 0; d < dim; d++) {
        float u_hat_d = u[d] + scale * w[d];
        float psi_d = h_prime * w[d];
        psi_dot_u += (double)psi_d * (double)u_hat_d;
        z_out[idx * dim + d] = z_in[idx * dim + d] + u_hat_d * h;
    }
    float abs_val = fabsf(1.0f + (float)psi_dot_u);
    float log_det = logf(abs_val + 1e-8f);
    sum_log_det[idx] += log_det;
}

std::tuple<torch::Tensor, torch::Tensor> planar_flow_forward(
    torch::Tensor z, 
    std::vector<torch::Tensor> ws, 
    std::vector<torch::Tensor> us, 
    std::vector<torch::Tensor> bs
) {
    int batch_size = z.size(0);
    int dim = z.size(1);
    int num_flows = ws.size();

    auto z_current = z.clone();
    auto sum_log_det = torch::zeros({batch_size}, z.options());

    const int threads = 256;
    const int blocks = (batch_size + threads - 1) / threads;

    for (int i = 0; i < num_flows; i++) {
        auto w = ws[i];
        auto u = us[i];
        auto b = bs[i];
        auto z_next = torch::empty_like(z_current);
        planar_flow_kernel<<<blocks, threads>>>(
            z_current.data_ptr<float>(),
            z_next.data_ptr<float>(),
            sum_log_det.data_ptr<float>(),
            w.data_ptr<float>(),
            u.data_ptr<float>(),
            b.data_ptr<float>(),
            batch_size,
            dim
        );
        z_current = z_next;
    }

    return std::make_tuple(z_current, sum_log_det);
}
"""

cpp_source = """
std::tuple<torch::Tensor, torch::Tensor> planar_flow_forward(
    torch::Tensor z, 
    std::vector<torch::Tensor> ws, 
    std::vector<torch::Tensor> us, 
    std::vector<torch::Tensor> bs
);
"""

module = load_inline(
    name="planar_flow",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["planar_flow_forward"],
    verbose=False,
    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)])
        self.module = module

    def forward(self, z):
        return self.module.planar_flow_forward(z, list(self.w), list(self.u), list(self.b))