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

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor affine_grid3d_cuda(torch::Tensor theta, int N, int D, int H, int W);
        """
        
        cuda_source = """
        #include <cuda_runtime.h>

        #define BLOCK_SIZE 256

        __global__ void affine_grid3d_packet_kernel(
            const float* __restrict__ theta,
            float* __restrict__ grid,
            int n_packets,
            int N, int D, int H, int W
        ) {
            int idx = blockIdx.x * blockDim.x + threadIdx.x;
            if (idx >= n_packets) return;

            int w_vec_dim = W / 4;
            
            int tmp = idx;
            int w_vec = tmp % w_vec_dim; tmp /= w_vec_dim;
            int h = tmp % H; tmp /= H;
            int d = tmp % D; 
            int n = tmp / D;

            int w_start = w_vec * 4;

         
            // Theta: [N, 3, 4] -> 12 floats
            const float* t = theta + n * 12;
            
      
            double t00=t[0], t01=t[1], t02=t[2], t03=t[3];
            double t10=t[4], t11=t[5], t12=t[6], t13=t[7];
            double t20=t[8], t21=t[9], t22=t[10], t23=t[11];

            double inv_w = 2.0 / (W - 1.0);
            double inv_h = 2.0 / (H - 1.0);
            double inv_d = 2.0 / (D - 1.0);

            double z = d * inv_d - 1.0;
            double y = h * inv_h - 1.0;

            float vals[12]; 

            #pragma unroll
            for (int i = 0; i < 4; ++i) {
                double x = (w_start + i) * inv_w - 1.0;
                
                // Affine Transform
                float px = (float)(t00 * x + t01 * y + t02 * z + t03);
                float py = (float)(t10 * x + t11 * y + t12 * z + t13);
                float pz = (float)(t20 * x + t21 * y + t22 * z + t23);
                
                vals[i*3 + 0] = px;
                vals[i*3 + 1] = py;
                vals[i*3 + 2] = pz;
            }

            float4* out_ptr = reinterpret_cast<float4*>(grid);
            long long out_idx = (long long)idx * 3;
            
            out_ptr[out_idx + 0] = ((float4*)vals)[0];
            out_ptr[out_idx + 1] = ((float4*)vals)[1];
            out_ptr[out_idx + 2] = ((float4*)vals)[2];
        }

        torch::Tensor affine_grid3d_cuda(torch::Tensor theta, int N, int D, int H, int W) {
            // Output: [N, D, H, W, 3]
            auto output = torch::empty({N, D, H, W, 3}, theta.options());
    
            if (W % 4 != 0) return output;

            long long total_voxels = (long long)N * D * H * W;
            int n_packets = total_voxels / 4;
            
            const int block_size = 256;
            const int grid_size = (n_packets + block_size - 1) / block_size;

            affine_grid3d_packet_kernel<<<grid_size, block_size>>>(
                theta.data_ptr<float>(),
                output.data_ptr<float>(),
                n_packets,
                N, D, H, W
            );
            
            return output;
        }
        """

        self.op = load_inline(
            name="affine_grid3d_packet_opt",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["affine_grid3d_cuda"],
            extra_cuda_cflags=["-O3"],
            verbose=False
        )

    def forward(self, theta: torch.Tensor) -> torch.Tensor:
        N = theta.shape[0]
        return self.op.affine_grid3d_cuda(theta, N, 64, 64, 64)