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 = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

__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 B,
    int dim,
    int num_flows
) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;

    float z[32];
    for(int d=0; d<dim; d++) {
        z[d] = z_in[b * dim + d];
    }
    float sld = 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];

        float wtu = 0.0f;
        for(int d=0; d<dim; d++) wtu += w[d] * u[d];
        float m = -1.0f + log1pf(expf(wtu));
        float w_norm2 = 0.0f;
        for(int d=0; d<dim; d++) w_norm2 += w[d]*w[d];
        float scale = (m - wtu) / w_norm2;
        float u_hat[32];
        for(int d=0; d<dim; d++) u_hat[d] = u[d] + scale * w[d];

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

        for(int d=0; d<dim; d++) {
            z[d] += u_hat[d] * h;
        }

        float psi_dot_u = 0.0f;
        for(int d=0; d<dim; d++) {
            psi_dot_u += h_prime * w[d] * u_hat[d];
        }
        float log_det = logf(fabsf(1.0f + psi_dot_u) + 1e-8f);
        sld += log_det;
    }

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

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 B = z.size(0);
    int dim = z.size(1);
    auto z_out = torch::empty_like(z);
    auto sld = torch::zeros({B}, z.options());

    dim3 blocks((B + 255)/256);
    dim3 threads(256);
    planar_flow_kernel<<<blocks, threads>>>(
        z.data_ptr<float>(),
        z_out.data_ptr<float>(),
        sld.data_ptr<float>(),
        w_all.data_ptr<float>(),
        u_all.data_ptr<float>(),
        b_all.data_ptr<float>(),
        B, dim, num_flows
    );
    return std::make_tuple(z_out, sld);
}
'''

cpp_source = r'''
#include <torch/extension.h>
#include <tuple>
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);
'''

_ext = load_inline(
    name="planar_flow",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["planar_flow_cuda"],
    verbose=False
)

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(list(self.w))
        u_all = torch.stack(list(self.u))
        b_all = torch.stack([bb.squeeze() for bb in self.b])
        return _ext.planar_flow_cuda(z, w_all, u_all, b_all, self.num_flows)