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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

cuda_source = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

__device__ __forceinline__ 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_all,
    const float* __restrict__ u_all,
    const float* __restrict__ b_all,
    int batch_size,
    int dim,
    int num_flows
) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch_size) return;

    float local_z[32];  // assume dim <= 32
    for (int d = 0; d < dim; d++) {
        local_z[d] = z_in[b * dim + d];
    }

    float log_det_acc = 0.0f;

    for (int i = 0; i < num_flows; i++) {
        const float* w = w_all + i * dim;
        const float* u = u_all + i * dim;
        float b = b_all[i];

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

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

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

    for (int d = 0; d < dim; d++) {
        z_out[b * dim + d] = local_z[d];
    }
    sum_log_det[b] = log_det_acc;
}

std::tuple<torch::Tensor, torch::Tensor> planar_flow_cuda(
    torch::Tensor z,
    torch::Tensor w_all,
    torch::Tensor u_all,
    torch::Tensor b_all,
    int num_flows
) {
    int batch_size = z.size(0);
    int dim = z.size(1);
    auto z_out = torch::empty_like(z);
    auto sum_log_det = torch::zeros({batch_size}, z.options());

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

    planar_flow_kernel<<<blocks, threads>>>(
        z.data_ptr<float>(),
        z_out.data_ptr<float>(),
        sum_log_det.data_ptr<float>(),
        w_all.data_ptr<float>(),
        u_all.data_ptr<float>(),
        b_all.data_ptr<float>(),
        batch_size,
        dim,
        num_flows
    );

    return std::make_tuple(z_out, sum_log_det);
}
'''

cpp_source = r'''
#include <torch/extension.h>
std::tuple<torch::Tensor, torch::Tensor> planar_flow_cuda(
    torch::Tensor z,
    torch::Tensor w_all,
    torch::Tensor u_all,
    torch::Tensor b_all,
    int num_flows
);
'''

module = load_inline(
    name='planar_flow_ext',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['planar_flow_cuda'],
    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)])

    def forward(self, z: torch.Tensor) -> tuple:
        w_all = torch.stack([w.data for w in self.w])
        u_all = torch.stack([u.data for u in self.u])
        b_all = torch.stack([b.data for b in self.b]).view(-1)
        return module.planar_flow_cuda(z, w_all, u_all, b_all, self.num_flows)
```