import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D convolution, scales the output, applies tanh, multiplies by a scaling factor, and applies sigmoid.
    """
    def __init__(self, in_channels, out_channels, kernel_size, scaling_factor, bias_shape):
        super(Model, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size)
        self.scaling_factor = nn.Parameter(torch.randn(bias_shape))
        self.bias = nn.Parameter(torch.randn(bias_shape)) 

    def forward(self, x):
        x = self.conv(x)
        x = x * self.scaling_factor 
        x = torch.tanh(x)
        x = x * self.bias
        x = torch.sigmoid(x)
        return x

batch_size = 128
in_channels = 3
out_channels = 16
depth, height, width = 16, 64, 64
kernel_size = 3
scaling_factor = 2
bias_shape = (out_channels, 1, 1, 1)

def get_inputs():
    return [torch.rand(batch_size, in_channels, depth, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, scaling_factor, bias_shape]

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

cpp_src = r'''
#include <torch/extension.h>

torch::Tensor conv3d_fused_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor scale, torch::Tensor bias);

torch::Tensor conv3d_fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor scale, torch::Tensor bias) {
    return conv3d_fused_cuda(x, weight, conv_bias, scale, bias);
}
'''

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

__global__ __launch_bounds__(128, 4)
void conv3d_fused_kernel(const float* __restrict__ x,
                         const float* __restrict__ w,
                         const float* __restrict__ cb,
                         const float* __restrict__ scale,
                         const float* __restrict__ bias,
                         float* __restrict__ out,
                         int total_spatial) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;

    constexpr int C = 3;
    constexpr int D = 16;
    constexpr int H = 64;
    constexpr int W = 64;
    constexpr int DO = 14;
    constexpr int HO = 62;
    constexpr int WO = 62;
    constexpr int OS = DO * HO * WO;
    constexpr int IS_C = D * H * W;
    constexpr int IS_N = C * IS_C;
    constexpr int OUT_N = 16 * OS;

    for (int idx = tid; idx < total_spatial; idx += stride) {
        int t = idx;
        int ox = t % WO; t /= WO;
        int oy = t % HO; t /= HO;
        int oz = t % DO; t /= DO;
        int n = t;

        float a0 = cb[0], a1 = cb[1], a2 = cb[2], a3 = cb[3];
        float a4 = cb[4], a5 = cb[5], a6 = cb[6], a7 = cb[7];
        float a8 = cb[8], a9 = cb[9], a10 = cb[10], a11 = cb[11];
        float a12 = cb[12], a13 = cb[13], a14 = cb[14], a15 = cb[15];

        int base_n = n * IS_N;
        #pragma unroll
        for (int c = 0; c < C; ++c) {
            int base_c = base_n + c * IS_C;
            #pragma unroll
            for (int kz = 0; kz < 3; ++kz) {
                int zbase = (oz + kz) * H * W;
                #pragma unroll
                for (int ky = 0; ky < 3; ++ky) {
                    int ybase = (oy + ky) * W;
                    #pragma unroll
                    for (int kx = 0; kx < 3; ++kx) {
                        int kidx = ((c * 3 + kz) * 3 + ky) * 3 + kx;
                        float xv = __ldg(x + base_c + zbase + ybase + ox + kx);
                        a0  = fmaf(xv, __ldg(w + 0  * 81 + kidx), a0);
                        a1  = fmaf(xv, __ldg(w + 1  * 81 + kidx), a1);
                        a2  = fmaf(xv, __ldg(w + 2  * 81 + kidx), a2);
                        a3  = fmaf(xv, __ldg(w + 3  * 81 + kidx), a3);
                        a4  = fmaf(xv, __ldg(w + 4  * 81 + kidx), a4);
                        a5  = fmaf(xv, __ldg(w + 5  * 81 + kidx), a5);
                        a6  = fmaf(xv, __ldg(w + 6  * 81 + kidx), a6);
                        a7  = fmaf(xv, __ldg(w + 7  * 81 + kidx), a7);
                        a8  = fmaf(xv, __ldg(w + 8  * 81 + kidx), a8);
                        a9  = fmaf(xv, __ldg(w + 9  * 81 + kidx), a9);
                        a10 = fmaf(xv, __ldg(w + 10 * 81 + kidx), a10);
                        a11 = fmaf(xv, __ldg(w + 11 * 81 + kidx), a11);
                        a12 = fmaf(xv, __ldg(w + 12 * 81 + kidx), a12);
                        a13 = fmaf(xv, __ldg(w + 13 * 81 + kidx), a13);
                        a14 = fmaf(xv, __ldg(w + 14 * 81 + kidx), a14);
                        a15 = fmaf(xv, __ldg(w + 15 * 81 + kidx), a15);
                    }
                }
            }
        }
        float s0=scale[0], s1=scale[1], s2=scale[2], s3=scale[3], s4=scale[4], s5=scale[5], s6=scale[6], s7=scale[7];
        float s8=scale[8], s9=scale[9], s10=scale[10], s11=scale[11], s12=scale[12], s13=scale[13], s14=scale[14], s15=scale[15];
        float b0=bias[0], b1=bias[1], b2=bias[2], b3=bias[3], b4=bias[4], b5=bias[5], b6=bias[6], b7=bias[7];
        float b8=bias[8], b9=bias[9], b10=bias[10], b11=bias[11], b12=bias[12], b13=bias[13], b14=bias[14], b15=bias[15];
        int obase = n * OUT_N + oz * HO * WO + oy * WO + ox;
        out[obase + 0*OS]  = 1.0f / (1.0f + expf(-(tanhf(a0  * s0)  * b0)));
        out[obase + 1*OS]  = 1.0f / (1.0f + expf(-(tanhf(a1  * s1)  * b1)));
        out[obase + 2*OS]  = 1.0f / (1.0f + expf(-(tanhf(a2  * s2)  * b2)));
        out[obase + 3*OS]  = 1.0f / (1.0f + expf(-(tanhf(a3  * s3)  * b3)));
        out[obase + 4*OS]  = 1.0f / (1.0f + expf(-(tanhf(a4  * s4)  * b4)));
        out[obase + 5*OS]  = 1.0f / (1.0f + expf(-(tanhf(a5  * s5)  * b5)));
        out[obase + 6*OS]  = 1.0f / (1.0f + expf(-(tanhf(a6  * s6)  * b6)));
        out[obase + 7*OS]  = 1.0f / (1.0f + expf(-(tanhf(a7  * s7)  * b7)));
        out[obase + 8*OS]  = 1.0f / (1.0f + expf(-(tanhf(a8  * s8)  * b8)));
        out[obase + 9*OS]  = 1.0f / (1.0f + expf(-(tanhf(a9  * s9)  * b9)));
        out[obase + 10*OS] = 1.0f / (1.0f + expf(-(tanhf(a10 * s10) * b10)));
        out[obase + 11*OS] = 1.0f / (1.0f + expf(-(tanhf(a11 * s11) * b11)));
        out[obase + 12*OS] = 1.0f / (1.0f + expf(-(tanhf(a12 * s12) * b12)));
        out[obase + 13*OS] = 1.0f / (1.0f + expf(-(tanhf(a13 * s13) * b13)));
        out[obase + 14*OS] = 1.0f / (1.0f + expf(-(tanhf(a14 * s14) * b14)));
        out[obase + 15*OS] = 1.0f / (1.0f + expf(-(tanhf(a15 * s15) * b15)));
    }
}

torch::Tensor conv3d_fused_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor scale, torch::Tensor bias) {
    auto out = torch::empty({x.size(0), weight.size(0), x.size(2)-2, x.size(3)-2, x.size(4)-2}, x.options());
    int total_spatial = x.size(0) * (x.size(2)-2) * (x.size(3)-2) * (x.size(4)-2);
    int threads = 128;
    int blocks = (total_spatial + threads - 1) / threads;
    if (blocks > 4096) blocks = 4096;
    conv3d_fused_kernel<<<blocks, threads>>>(x.data_ptr<float>(), weight.data_ptr<float>(), conv_bias.data_ptr<float>(), scale.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(), total_spatial);
    return out;
}
'''

_ext = load_inline(name='fused_conv3d_act_ext_submit', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['conv3d_fused_forward'], extra_cuda_cflags=['-O3', '--use_fast_math'], extra_cflags=['-O3'], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, scaling_factor, bias_shape):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size)
        self.scaling_factor = nn.Parameter(torch.randn(bias_shape))
        self.bias = nn.Parameter(torch.randn(bias_shape))
    def forward(self, x):
        return _ext.conv3d_fused_forward(x.contiguous(), self.conv.weight.contiguous(), self.conv.bias.contiguous(), self.scaling_factor.contiguous(), self.bias.contiguous())