#include <torch/extension.h>
#include <cuda_runtime.h>
#include <vector>

#define BM 64
#define BN 128
#define BK 8
#define THREADS 128

__global__ void __launch_bounds__(THREADS)
conv_mish_kernel(const float* __restrict__ x,
                 const float* __restrict__ w,
                 const float* __restrict__ bias,
                 float* __restrict__ out,
                 int Cin, int H, int W, int Cout, int K,
                 int Hout, int Wout, long Nflat,
                 float s1, float s2) {
    __shared__ float As[2][BK][BN];
    __shared__ float Bs[2][BM][BK + 1];
    __shared__ int e0s[BN];
    __shared__ long long outs[BN];
    __shared__ int koff[72];

    const int tid = threadIdx.x;
    const int ty = tid >> 4;   // 0..7  (m dimension)
    const int tx = tid & 15;   // 0..15 (n dimension)
    const int HW = H * W;
    const int Ktot = Cin * K * K;
    const int nchunk = (Hout * Wout + BN - 1) / BN;

    const long nblk = blockIdx.x % nchunk;
    const int img = (int)(blockIdx.x / nchunk);
    const long n0 = nblk * BN;
    const long imgOff = (long)img * HW;

    // Precompute flattened-index bookkeeping
    if (tid < 72) {
        if (tid < Ktot) {
            int ic = tid / (K * K);
            int r = tid % (K * K);
            koff[tid] = ic * HW + (r / K) * W + (r % K);
        } else {
            koff[tid] = 0;
        }
    }
    {
        long bn = n0 + tid;
        if (bn < Nflat) {
            int p = (int)(bn / Wout);
            int q = (int)(bn % Wout);
            e0s[tid] = p * W + q;
            outs[tid] = (long)p * Wout + q;
        } else {
            e0s[tid] = 0;
            outs[tid] = -1;
        }
    }
    __syncthreads();

    float acc[8][8];
#pragma unroll
    for (int i = 0; i < 8; ++i) {
        float b = (ty * 8 + i) < Cout ? bias[ty * 8 + i] : 0.f;
#pragma unroll
        for (int j = 0; j < 8; ++j) acc[i][j] = b;
    }

    float pa[BK], pb[4];
    auto load_chunk = [&](int c, float* ppa, float* ppb) {
        int kb = c * BK;
#pragma unroll
        for (int i = 0; i < BK; ++i) {
            int k = kb + i;
            int off = (k < Ktot) ? koff[k] : 0;
            ppa[i] = __ldg(x + imgOff + off + e0s[tx * 8 / 8 * 8 + (i * THREADS + tid) % BN]); // placeholder
        }
    };
    // (lambda above unused; explicit code below)

    // ---- prefetch chunk 0 ----
#pragma unroll
    for (int i = 0; i < BK; ++i) {
        int idx = i * THREADS + tid;      // k = idx / BN, n = idx % BN
        int k = idx >> 7;
        int n = idx & 127;
        pa[i] = __ldg(x + imgOff + koff[k] + e0s[n]);
    }
#pragma unroll
    for (int i = 0; i < 4; ++i) {
        int idx = i * THREADS + tid;      // m = idx / BK, kk = idx % BK
        int m = idx >> 3;
        int kk = idx & 7;
        pb[i] = (m < Cout && kk < Ktot) ? __ldg(w + (long)m * Ktot + kk) : 0.f;
    }
#pragma unroll
    for (int i = 0; i < BK; ++i) {
        int idx = i * THREADS + tid;
        As[0][idx >> 7][idx & 127] = pa[i];
    }
#pragma unroll
    for (int i = 0; i < 4; ++i) {
        int idx = i * THREADS + tid;
        Bs[0][idx >> 3][idx & 7] = pb[i];
    }
    __syncthreads();

    const int nchunks = (Ktot + BK - 1) / BK;
    for (int c = 0; c < nchunks; ++c) {
        int cur = c & 1;
        if (c + 1 < nchunks) {
            int kb = (c + 1) * BK;
#pragma unroll
            for (int i = 0; i < BK; ++i) {
                int idx = i * THREADS + tid;
                int k = kb + (idx >> 7);
                int n = idx & 127;
                pa[i] = (k < Ktot) ? __ldg(x + imgOff + koff[k] + e0s[n]) : 0.f;
            }
#pragma unroll
            for (int i = 0; i < 4; ++i) {
                int idx = i * THREADS + tid;
                int m = idx >> 3;
                int kk = kb + (idx & 7);
                pb[i] = (m < Cout && kk < Ktot) ? __ldg(w + (long)m * Ktot + kk) : 0.f;
            }
        }
#pragma unroll
        for (int kk = 0; kk < BK; ++kk) {
            float af[8], bf[8];
            const float4* aptr = reinterpret_cast<const float4*>(&As[cur][kk][tx * 8]);
            float4 a0 = aptr[0], a1 = aptr[1];
            af[0] = a0.x; af[1] = a0.y; af[2] = a0.z; af[3] = a0.w;
            af[4] = a1.x; af[5] = a1.y; af[6] = a1.z; af[7] = a1.w;
#pragma unroll
            for (int i = 0; i < 8; ++i) bf[i] = Bs[cur][ty * 8 + i][kk];
#pragma unroll
            for (int i = 0; i < 8; ++i)
#pragma unroll
                for (int j = 0; j < 8; ++j) acc[i][j] = fmaf(bf[i], af[j], acc[i][j]);
        }
        if (c + 1 < nchunks) {
            int nxt = cur ^ 1;
#pragma unroll
            for (int i = 0; i < BK; ++i) {
                int idx = i * THREADS + tid;
                As[nxt][idx >> 7][idx & 127] = pa[i];
            }
#pragma unroll
            for (int i = 0; i < 4; ++i) {
                int idx = i * THREADS + tid;
                Bs[nxt][idx >> 3][idx & 7] = pb[i];
            }
        }
        __syncthreads();
    }

    // Epilogue: subtract s1, subtract s2, Mish, store
#pragma unroll
    for (int j = 0; j < 8; ++j) {
        int n = tx * 8 + j;
        long ob = outs[n];
        if (ob >= 0) {
#pragma unroll
            for (int i = 0; i < 8; ++i) {
                int m = ty * 8 + i;
                if (m < Cout) {
                    float v = acc[i][j] - s1;
                    v = v - s2;
                    float sp = (v > 20.f) ? v : log1pf(expf(v));
                    v = v * tanhf(sp);
                    out[ob + (long)m * ((long)Hout * Wout)] = v;
                }
            }
        }
    }
}

torch::Tensor conv_mish_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b,
                             double s1, double s2) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda() && b.is_cuda(), "CUDA tensors required");
    auto xc = x.contiguous();
    auto wc = w.contiguous();
    auto bc = b.contiguous();

    const int N = xc.size(0);
    const int Cin = xc.size(1);
    const int H = xc.size(2);
    const int W = xc.size(3);
    const int Cout = wc.size(0);
    const int K = wc.size(2);
    const int Hout = H - K + 1;
    const int Wout = W - K + 1;
    const long Nflat = (long)N * Hout * Wout;

    auto out = torch::empty({N, Cout, Hout, Wout}, xc.options());

    const int nchunk = (Hout * Wout + BN - 1) / BN;
    dim3 grid((unsigned)((long)nchunk * N));
    conv_mish_kernel<<<grid, THREADS>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), bc.data_ptr<float>(),
        out.data_ptr<float>(), Cin, H, W, Cout, K, Hout, Wout, Nflat,
        (float)s1, (float)s2);
    return out;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("conv_mish", &conv_mish_cuda, "Fused conv2d + sub + sub + mish (CUDA)");
}
