#include "pipeline.h"

/* Structure tensor from Jähne's rotationally optimized Sobel derivatives */

__global__ void structure_tensor(
    const float4* src, float4* dst, int w, int h)
{
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;
    if (x >= w || y >= h) return;

    const float p1 = 0.183f;
    const float p2 = 1.0f - 2.0f * p1;

    int xm = max(x - 1, 0), xp = min(x + 1, w - 1);
    int ym = max(y - 1, 0), yp = min(y + 1, h - 1);

    float4 c_mm = src[ym * w + xm], c_0m = src[ym * w + x], c_pm = src[ym * w + xp];
    float4 c_m0 = src[y  * w + xm],                         c_p0 = src[y  * w + xp];
    float4 c_mp = src[yp * w + xm], c_0p = src[yp * w + x], c_pp = src[yp * w + xp];

    float3 gx, gy;
    gx.x = 0.5f * (p1 * (c_pm.x - c_mm.x) + p2 * (c_p0.x - c_m0.x) + p1 * (c_pp.x - c_mp.x));
    gx.y = 0.5f * (p1 * (c_pm.y - c_mm.y) + p2 * (c_p0.y - c_m0.y) + p1 * (c_pp.y - c_mp.y));
    gx.z = 0.5f * (p1 * (c_pm.z - c_mm.z) + p2 * (c_p0.z - c_m0.z) + p1 * (c_pp.z - c_mp.z));
    gy.x = 0.5f * (p1 * (c_mp.x - c_mm.x) + p2 * (c_0p.x - c_0m.x) + p1 * (c_pp.x - c_pm.x));
    gy.y = 0.5f * (p1 * (c_mp.y - c_mm.y) + p2 * (c_0p.y - c_0m.y) + p1 * (c_pp.y - c_pm.y));
    gy.z = 0.5f * (p1 * (c_mp.z - c_mm.z) + p2 * (c_0p.z - c_0m.z) + p1 * (c_pp.z - c_pm.z));

    dst[y * w + x] = make_float4(
        gx.x * gx.x + gx.y * gx.y + gx.z * gx.z,
        gx.x * gy.x + gx.y * gy.y + gx.z * gy.z,
        gy.x * gy.x + gy.y * gy.y + gy.z * gy.z,
        0.0f);
}

__device__ float4 gaussian_sample_axis(
    const float4* src, int src_len, int src_stride, int src_base,
    int center, float sigma)
{
    int radius = static_cast<int>(ceilf(3.0f * sigma));
    float inv_two_sigma2 = 1.0f / (2.0f * sigma * sigma);

    float4 sum = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
    float wsum = 0.0f;

    for (int i = -radius; i <= radius; ++i) {
        int c = min(max(center + i, 0), src_len - 1);
        float wt = expf(-(i * i) * inv_two_sigma2);
        float4 s = src[src_base + c * src_stride];
        sum.x += s.x * wt; sum.y += s.y * wt;
        sum.z += s.z * wt; sum.w += s.w * wt;
        wsum += wt;
    }

    float inv = 1.0f / wsum;
    sum.x *= inv; sum.y *= inv; sum.z *= inv; sum.w *= inv;
    return sum;
}

__global__ void gaussian_horizontal(
    const float4* src, float4* dst, int w, int h, float sigma)
{
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;
    if (x >= w || y >= h) return;

    dst[y * w + x] = gaussian_sample_axis(src, w, 1, y * w, x, sigma);
}

__global__ void gaussian_vertical(
    const float4* src, float4* dst, int w, int h, float sigma)
{
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;
    if (x >= w || y >= h) return;

    dst[y * w + x] = gaussian_sample_axis(src, h, w, x, y, sigma);
}

void compute_tensor(PyramidLevel& lvl) {
    dim3 block(16, 16);
    dim3 grid((lvl.width + 15) / 16, (lvl.height + 15) / 16);

    structure_tensor<<<grid, block>>>(
        lvl.color_smax, lvl.tensor, lvl.width, lvl.height);
    hip_check(hipGetLastError());
}

void smooth_tensor(PyramidLevel& lvl, const Params& p) {
    float4* tmp;
    hip_check(hipMalloc(&tmp, sizeof(float4) * lvl.width * lvl.height));

    dim3 block(16, 16);
    dim3 grid((lvl.width + 15) / 16, (lvl.height + 15) / 16);

    gaussian_horizontal<<<grid, block>>>(
        lvl.tensor, tmp, lvl.width, lvl.height, p.rho);
    hip_check(hipGetLastError());

    gaussian_vertical<<<grid, block>>>(
        tmp, lvl.tensor, lvl.width, lvl.height, p.rho);
    hip_check(hipGetLastError());

    hipFree(tmp);
}