import torch
from torch.utils.cpp_extension import load_inline

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

#define N 8

// ---------------------------------------------------------
// Device Function: In-Register Cholesky Decomposition
// Computes LogDet directly
// ---------------------------------------------------------
__device__ __forceinline__ float compute_logdet_8x8(const float* mat_in) {
    // 1. Load Matrix into Registers
    float A[N][N];
    
    #pragma unroll
    for (int i = 0; i < N; ++i) {
        #pragma unroll
        for (int j = 0; j < N; ++j) {
            A[i][j] = mat_in[i * N + j];
        }
    }
    
    // 2. Cholesky Decomposition (L * L^T = A)
    // We only need diagonal elements L_ii to compute LogDet
    // In-place update: A will store L (lower triangular)
    
    double logdet_sum = 0.0;
    
    // Unrolled Cholesky
    #pragma unroll
    for (int k = 0; k < N; ++k) {
        
        // Compute L_kk
        float diag_val = A[k][k];
        
        // diag_val -= sum(L_kj ^ 2) for j < k
        for (int j = 0; j < k; ++j) {
            float val = A[k][j];
            diag_val -= val * val;
        }
        
        // Check Positive Definite
        if (diag_val <= 0.0f) return -1e9f; // Error case (should not happen for SPD)
        
        float L_kk = sqrtf(diag_val);
        A[k][k] = L_kk;
        
        // Accumulate LogDet: 2 * sum(ln(L_ii))
        logdet_sum += log((double)L_kk);
        
        // Compute L_ik for i > k
        float inv_Lkk = 1.0f / L_kk;
        
        for (int i = k + 1; i < N; ++i) {
            float val = A[i][k];
            
            // val -= sum(L_ij * L_kj) for j < k
            for (int j = 0; j < k; ++j) {
                val -= A[i][j] * A[k][j];
            }
            
            A[i][k] = val * inv_Lkk;
        }
    }
    
    return (float)(2.0 * logdet_sum);
}

__global__ void batched_logdet_8x8_kernel(
    const float* __restrict__ input, 
    float* __restrict__ output, 
    int batch_size) 
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (idx < batch_size) {
        // Point to the current matrix
        const float* mat_ptr = input + idx * (N * N);
        
        // Compute
        float res = compute_logdet_8x8(mat_ptr);
        
        // Write Output
        output[idx] = res;
    }
}

torch::Tensor logdet_cuda(torch::Tensor input) {
    int B = input.size(0);
    // Input check: input.size(1) == 8 && input.size(2) == 8
    
    auto output = torch::empty({B}, input.options());
    
    const int block = 256;
    const int grid = (B + block - 1) / block;
    
    batched_logdet_8x8_kernel<<<grid, block>>>(
        input.data_ptr<float>(),
        output.data_ptr<float>(),
        B
    );
    
    return output;
}
"""

cpp_source = "torch::Tensor logdet_cuda(torch::Tensor input);"

logdet_module = load_inline(
    name="logdet_extension",
    cpp_sources=cpp_source,
    cuda_sources=logdet_source,
    functions=["logdet_cuda"],
    verbose=True,
    with_cuda=True
)

class ModelNew(torch.nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
        self.cuda_op = logdet_module

    def forward(self, x):
        # 确保输入连续
        return self.cuda_op.logdet_cuda(x.contiguous())