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

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

// ---------------------------------------------------------
// Helper: 4x4 Matrix Multiplication (Registers)
// Using double for accumulator to ensure precision
// ---------------------------------------------------------
__device__ __forceinline__ void matmul_4x4(
    const float A[16], 
    const float B[16], 
    float C[16]) 
{
    // Unroll loops manually
    #pragma unroll
    for (int i = 0; i < 4; ++i) {
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            double sum = 0.0; // Use double accumulator
            #pragma unroll
            for (int k = 0; k < 4; ++k) {
                sum += (double)A[i*4 + k] * (double)B[k*4 + j];
            }
            C[i*4 + j] = (float)sum;
        }
    }
}

// ---------------------------------------------------------
// Helper: Copy 16 floats
// ---------------------------------------------------------
__device__ __forceinline__ void copy_mat(const float src[16], float dst[16]) {
    #pragma unroll
    for(int i=0; i<16; ++i) dst[i] = src[i];
}

// ---------------------------------------------------------
// Kernel: Batched Matrix Power
// ---------------------------------------------------------
__global__ void matrix_power_4x4_kernel(
    const float* __restrict__ input, 
    float* __restrict__ output, 
    int batch_size, 
    int power) 
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (idx < batch_size) {
        // Pointers
        const float* in_ptr = input + idx * 16;
        float* out_ptr = output + idx * 16;
        
        // Registers
        float base[16];
        float res[16];
        float tmp[16];
        
        // Load Input (Vectorized Float4)
        const float4* v_in = reinterpret_cast<const float4*>(in_ptr);
        float4* v_base = reinterpret_cast<float4*>(base);
        
        v_base[0] = v_in[0];
        v_base[1] = v_in[1];
        v_base[2] = v_in[2];
        v_base[3] = v_in[3];

        // Initialize Res to Identity
        #pragma unroll
        for(int i=0; i<16; ++i) res[i] = 0.0f;
        res[0] = 1.0f; res[5] = 1.0f; res[10] = 1.0f; res[15] = 1.0f;
        
        // Binary Exponentiation
        int p = power;
        
        while (p > 0) {
            if (p & 1) {
                matmul_4x4(res, base, tmp);
                copy_mat(tmp, res);
            }
            
            p >>= 1;
            
            if (p > 0) {
                matmul_4x4(base, base, tmp);
                copy_mat(tmp, base);
            }
        }

        // Store Result (Vectorized)
        float4* v_out = reinterpret_cast<float4*>(out_ptr);
        float4* v_res = reinterpret_cast<float4*>(res);
        
        v_out[0] = v_res[0];
        v_out[1] = v_res[1];
        v_out[2] = v_res[2];
        v_out[3] = v_res[3];
    }
}

torch::Tensor matrix_power_cuda(torch::Tensor input, int power) {
    int B = input.size(0);
    auto output = torch::empty_like(input);
    
    const int block = 256;
    const int grid = (B + block - 1) / block;
    
    matrix_power_4x4_kernel<<<grid, block>>>(
        input.data_ptr<float>(),
        output.data_ptr<float>(),
        B,
        power
    );
    
    return output;
}
"""

cpp_source = "torch::Tensor matrix_power_cuda(torch::Tensor input, int power);"

mpow_module = load_inline(
    name="matrix_power_extension_v2",
    cpp_sources=cpp_source,
    cuda_sources=mpow_source,
    functions=["matrix_power_cuda"],
    verbose=True,
    with_cuda=True
)

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

    def forward(self, x):
        return self.cuda_op.matrix_power_cuda(x.contiguous(), self.power)